747 lines
22 KiB
Python
747 lines
22 KiB
Python
import os
|
|
import random
|
|
import importlib
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
try:
|
|
import comfy.samplers
|
|
except ImportError:
|
|
comfy = None
|
|
|
|
|
|
_COMMON_KSAMPLER_OVERRIDE = None
|
|
AXIS_OPTIONS = ["seed", "steps", "cfg", "denoise"]
|
|
|
|
|
|
def hex_to_rgb(hex_color: str):
|
|
hex_color = hex_color.strip().lstrip("#")
|
|
|
|
if len(hex_color) != 6:
|
|
raise ValueError("Color must be a hex string like #FFFFFF")
|
|
|
|
return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
|
|
|
|
|
|
def tensor_to_pil(image_tensor):
|
|
"""
|
|
ComfyUI IMAGE tensor: [H, W, C], range 0~1
|
|
"""
|
|
image_np = image_tensor.cpu().numpy()
|
|
image_np = np.clip(image_np * 255.0, 0, 255).astype(np.uint8)
|
|
return Image.fromarray(image_np)
|
|
|
|
|
|
def pil_to_tensor(image):
|
|
"""
|
|
Return ComfyUI IMAGE tensor: [H, W, C], range 0~1
|
|
"""
|
|
image_np = np.asarray(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image_np)
|
|
|
|
|
|
def get_text_size(draw, text, font):
|
|
if hasattr(draw, "textbbox"):
|
|
left, top, right, bottom = draw.textbbox((0, 0), text, font=font)
|
|
return right - left, bottom - top
|
|
|
|
return draw.textsize(text, font=font)
|
|
|
|
|
|
def calculate_text_x(align, image_width, text_width, x):
|
|
if align == "center":
|
|
return int((image_width - text_width) / 2) + int(x)
|
|
if align == "right":
|
|
return int(image_width) - int(text_width) - int(x)
|
|
return int(x)
|
|
|
|
|
|
class StatusBarOverlay:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"status": (
|
|
"STRING",
|
|
{
|
|
"default": "status",
|
|
"multiline": False,
|
|
},
|
|
),
|
|
"bar_height": (
|
|
"INT",
|
|
{
|
|
"default": 30,
|
|
"min": 1,
|
|
"max": 512,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"bar_color": (
|
|
"STRING",
|
|
{
|
|
"default": "#FFFFFF",
|
|
},
|
|
),
|
|
"text_color": (
|
|
"STRING",
|
|
{
|
|
"default": "#000000",
|
|
},
|
|
),
|
|
"font_size": (
|
|
"INT",
|
|
{
|
|
"default": 18,
|
|
"min": 1,
|
|
"max": 256,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"x": (
|
|
"INT",
|
|
{
|
|
"default": 8,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"y": (
|
|
"INT",
|
|
{
|
|
"default": 5,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"align": (
|
|
["left", "center", "right"],
|
|
{
|
|
"default": "left",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"font_path": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "overlay_status_bar"
|
|
CATEGORY = "image/overlay"
|
|
|
|
def overlay_status_bar(
|
|
self,
|
|
image,
|
|
status,
|
|
bar_height,
|
|
bar_color,
|
|
text_color,
|
|
font_size,
|
|
x,
|
|
y,
|
|
align,
|
|
font_path="",
|
|
):
|
|
bar_rgb = hex_to_rgb(bar_color)
|
|
text_rgb = hex_to_rgb(text_color)
|
|
|
|
output_images = []
|
|
|
|
for img in image:
|
|
pil_img = tensor_to_pil(img).convert("RGB")
|
|
draw = ImageDraw.Draw(pil_img)
|
|
|
|
width, height = pil_img.size
|
|
safe_bar_height = min(int(bar_height), height)
|
|
|
|
draw.rectangle(
|
|
[(0, 0), (width - 1, safe_bar_height - 1)],
|
|
fill=bar_rgb,
|
|
)
|
|
|
|
if font_path and os.path.exists(font_path):
|
|
try:
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
else:
|
|
try:
|
|
font = ImageFont.truetype("arial.ttf", font_size)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
|
|
text_width, _ = get_text_size(draw, status, font)
|
|
text_x = calculate_text_x(align, width, text_width, x)
|
|
|
|
draw.text(
|
|
(text_x, y),
|
|
status,
|
|
fill=text_rgb,
|
|
font=font,
|
|
)
|
|
|
|
output_images.append(pil_to_tensor(pil_img))
|
|
|
|
return (torch.stack(output_images, dim=0),)
|
|
|
|
|
|
def _sampler_names():
|
|
if comfy is not None:
|
|
return comfy.samplers.KSampler.SAMPLERS
|
|
return ["euler"]
|
|
|
|
|
|
def _scheduler_names():
|
|
if comfy is not None:
|
|
return comfy.samplers.KSampler.SCHEDULERS
|
|
return ["normal"]
|
|
|
|
|
|
def _get_common_ksampler():
|
|
if _COMMON_KSAMPLER_OVERRIDE is not None:
|
|
return _COMMON_KSAMPLER_OVERRIDE
|
|
|
|
comfy_nodes = importlib.import_module("nodes")
|
|
if not hasattr(comfy_nodes, "common_ksampler"):
|
|
raise RuntimeError("ComfyUI common_ksampler was not found.")
|
|
return comfy_nodes.common_ksampler
|
|
|
|
|
|
def _random_seeds(start, end, count, generator_seed):
|
|
if start > end:
|
|
raise ValueError("seed_min must be less than or equal to seed_max")
|
|
rng = random.Random(int(generator_seed))
|
|
return [rng.randint(start, end) for _ in range(count)]
|
|
|
|
|
|
def _linspace_int(start, end, count, error_name):
|
|
if start > end:
|
|
raise ValueError(f"{error_name}_min must be less than or equal to {error_name}_max")
|
|
if count == 1:
|
|
return [int(round(start))]
|
|
|
|
step = (end - start) / (count - 1)
|
|
return [int(round(start + (index * step))) for index in range(count)]
|
|
|
|
|
|
def _linspace_float(start, end, count, decimals):
|
|
if start > end:
|
|
raise ValueError("axis min must be less than or equal to axis max")
|
|
if count == 1:
|
|
return [round(float(start), decimals)]
|
|
|
|
step = (end - start) / (count - 1)
|
|
return [round(start + (index * step), decimals) for index in range(count)]
|
|
|
|
|
|
def _parse_int_sequence(value, name, expected_count):
|
|
values = [int(part.strip()) for part in str(value).split(",") if part.strip()]
|
|
if not values:
|
|
raise ValueError(f"{name} must contain at least one integer.")
|
|
if len(values) not in (1, expected_count):
|
|
raise ValueError(f"{name} must contain either 1 value or {expected_count} values.")
|
|
return values
|
|
|
|
|
|
def _parse_float_sequence(value, name, expected_count):
|
|
values = [float(part.strip()) for part in str(value).split(",") if part.strip()]
|
|
if not values:
|
|
raise ValueError(f"{name} must contain at least one float.")
|
|
if len(values) not in (1, expected_count):
|
|
raise ValueError(f"{name} must contain either 1 value or {expected_count} values.")
|
|
return values
|
|
|
|
|
|
def _axis_values(axis, count, min_value, max_value, seed_generator_seed):
|
|
count = int(count)
|
|
if count < 1 or count > 5:
|
|
raise ValueError("axis count must be between 1 and 5")
|
|
|
|
if axis == "seed":
|
|
return _random_seeds(
|
|
int(round(min_value)),
|
|
int(round(max_value)),
|
|
count,
|
|
seed_generator_seed,
|
|
)
|
|
if axis == "steps":
|
|
return _linspace_int(float(min_value), float(max_value), count, "steps")
|
|
if axis == "cfg":
|
|
return _linspace_float(float(min_value), float(max_value), count, 2)
|
|
if axis == "denoise":
|
|
return _linspace_float(float(min_value), float(max_value), count, 3)
|
|
raise ValueError(f"Unsupported axis: {axis}")
|
|
|
|
|
|
def _apply_axis_value(params, axis, value):
|
|
if axis == "seed":
|
|
params["seed"] = int(value)
|
|
elif axis == "steps":
|
|
params["steps"] = int(value)
|
|
elif axis == "cfg":
|
|
params["cfg"] = float(value)
|
|
elif axis == "denoise":
|
|
params["denoise"] = float(value)
|
|
else:
|
|
raise ValueError(f"Unsupported axis: {axis}")
|
|
|
|
|
|
def _format_run_status(params):
|
|
return (
|
|
f"seed={params['seed']}, steps={params['steps']}, cfg={params['cfg']}, "
|
|
f"denoise={params['denoise']}, sampler={params['sampler_name']}, scheduler={params['scheduler']}"
|
|
)
|
|
|
|
|
|
def _format_changed_values(params, x_axis, y_axis):
|
|
return f"{x_axis}={params[x_axis]}, {y_axis}={params[y_axis]}"
|
|
|
|
|
|
class KSamplerMultiRun:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"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,
|
|
"round": 0.01,
|
|
},
|
|
),
|
|
"sampler_name": (_sampler_names(),),
|
|
"scheduler": (_scheduler_names(),),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"denoise": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
},
|
|
),
|
|
"seed_generator_seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"x_axis": (
|
|
AXIS_OPTIONS,
|
|
{
|
|
"default": "seed",
|
|
},
|
|
),
|
|
"x_count": (
|
|
"INT",
|
|
{
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 5,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"x_min": (
|
|
"FLOAT",
|
|
{
|
|
"default": 100000000.0,
|
|
"min": -999999999.0,
|
|
"max": 999999999.0,
|
|
"step": 0.01,
|
|
},
|
|
),
|
|
"x_max": (
|
|
"FLOAT",
|
|
{
|
|
"default": 999999999.0,
|
|
"min": -999999999.0,
|
|
"max": 999999999.0,
|
|
"step": 0.01,
|
|
},
|
|
),
|
|
"y_axis": (
|
|
AXIS_OPTIONS,
|
|
{
|
|
"default": "cfg",
|
|
},
|
|
),
|
|
"y_count": (
|
|
"INT",
|
|
{
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 5,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"y_min": (
|
|
"FLOAT",
|
|
{
|
|
"default": 4.0,
|
|
"min": -999999999.0,
|
|
"max": 999999999.0,
|
|
"step": 0.01,
|
|
},
|
|
),
|
|
"y_max": (
|
|
"FLOAT",
|
|
{
|
|
"default": 8.0,
|
|
"min": -999999999.0,
|
|
"max": 999999999.0,
|
|
"step": 0.01,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "STRING", "VALUE_LIST", "VALUE_LIST", "INT", "INT", "INT")
|
|
RETURN_NAMES = (
|
|
"samples",
|
|
"run_info",
|
|
"status_list",
|
|
"changed_values_list",
|
|
"columns",
|
|
"rows",
|
|
"batch_size",
|
|
)
|
|
FUNCTION = "sample_multi"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample_multi(
|
|
self,
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise,
|
|
seed_generator_seed,
|
|
x_axis,
|
|
x_count,
|
|
x_min,
|
|
x_max,
|
|
y_axis,
|
|
y_count,
|
|
y_min,
|
|
y_max,
|
|
):
|
|
if x_axis == y_axis:
|
|
raise ValueError("x_axis and y_axis must be different.")
|
|
|
|
x_count = int(x_count)
|
|
y_count = int(y_count)
|
|
run_count = x_count * y_count
|
|
|
|
if x_count > 5 or y_count > 5 or run_count > 25:
|
|
raise ValueError("KSampler Multi Run is limited to 5 x 5, max 25 runs.")
|
|
|
|
x_values = _axis_values(
|
|
x_axis,
|
|
x_count,
|
|
float(x_min),
|
|
float(x_max),
|
|
int(seed_generator_seed),
|
|
)
|
|
y_values = _axis_values(
|
|
y_axis,
|
|
y_count,
|
|
float(y_min),
|
|
float(y_max),
|
|
int(seed_generator_seed),
|
|
)
|
|
common_ksampler = _get_common_ksampler()
|
|
|
|
output_samples = []
|
|
run_lines = []
|
|
status_list = []
|
|
changed_values_list = []
|
|
|
|
for y_index, y_value in enumerate(y_values):
|
|
for x_index, x_value in enumerate(x_values):
|
|
run_index = (y_index * x_count) + x_index
|
|
params = {
|
|
"seed": int(seed),
|
|
"steps": int(steps),
|
|
"cfg": float(cfg),
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
"denoise": float(denoise),
|
|
}
|
|
_apply_axis_value(params, x_axis, x_value)
|
|
_apply_axis_value(params, y_axis, y_value)
|
|
|
|
sampled = common_ksampler(
|
|
model,
|
|
params["seed"],
|
|
params["steps"],
|
|
params["cfg"],
|
|
params["sampler_name"],
|
|
params["scheduler"],
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=params["denoise"],
|
|
)[0]
|
|
status = _format_run_status(params)
|
|
changed_values = _format_changed_values(params, x_axis, y_axis)
|
|
output_samples.append(sampled["samples"])
|
|
status_list.append(status)
|
|
changed_values_list.append(changed_values)
|
|
run_lines.append(f"{run_index}: {status}")
|
|
|
|
samples = torch.cat(output_samples, dim=0)
|
|
output_latent = dict(latent_image)
|
|
output_latent["samples"] = samples
|
|
output_latent.pop("batch_index", None)
|
|
run_info = (
|
|
f"runs={run_count}, columns={x_count}, rows={y_count}, "
|
|
f"x_axis={x_axis}, y_axis={y_axis}, seed_generator_seed={seed_generator_seed}\n"
|
|
+ "\n".join(run_lines)
|
|
)
|
|
|
|
return (output_latent, run_info, status_list, changed_values_list, x_count, y_count, run_count)
|
|
|
|
|
|
class BatchStatusBarOverlay:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"status_list": ("VALUE_LIST",),
|
|
"bar_height": ("INT", {"default": 30, "min": 1, "max": 512, "step": 1}),
|
|
"bar_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"text_color": ("STRING", {"default": "#000000"}),
|
|
"font_size": ("INT", {"default": 18, "min": 1, "max": 256, "step": 1}),
|
|
"x": ("INT", {"default": 8, "min": 0, "max": 4096, "step": 1}),
|
|
"y": ("INT", {"default": 5, "min": 0, "max": 4096, "step": 1}),
|
|
"align": (["left", "center", "right"], {"default": "left"}),
|
|
},
|
|
"optional": {
|
|
"font_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "overlay_status_list"
|
|
CATEGORY = "image/overlay"
|
|
|
|
def overlay_status_list(
|
|
self,
|
|
image,
|
|
status_list,
|
|
bar_height,
|
|
bar_color,
|
|
text_color,
|
|
font_size,
|
|
x,
|
|
y,
|
|
align,
|
|
font_path="",
|
|
):
|
|
statuses = list(status_list)
|
|
if len(statuses) < len(image):
|
|
raise ValueError(f"status_list has {len(statuses)} values, but image batch has {len(image)} images.")
|
|
|
|
bar_rgb = hex_to_rgb(bar_color)
|
|
text_rgb = hex_to_rgb(text_color)
|
|
font = None
|
|
output_images = []
|
|
|
|
for index, img in enumerate(image):
|
|
pil_img = tensor_to_pil(img).convert("RGB")
|
|
draw = ImageDraw.Draw(pil_img)
|
|
width, height = pil_img.size
|
|
safe_bar_height = min(int(bar_height), height)
|
|
draw.rectangle([(0, 0), (width - 1, safe_bar_height - 1)], fill=bar_rgb)
|
|
|
|
if font is None:
|
|
if font_path and os.path.exists(font_path):
|
|
try:
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
else:
|
|
try:
|
|
font = ImageFont.truetype("arial.ttf", font_size)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
|
|
status = str(statuses[index])
|
|
text_width, _ = get_text_size(draw, status, font)
|
|
text_x = calculate_text_x(align, width, text_width, x)
|
|
draw.text((text_x, y), status, fill=text_rgb, font=font)
|
|
output_images.append(pil_to_tensor(pil_img))
|
|
|
|
return (torch.stack(output_images, dim=0),)
|
|
|
|
|
|
def hex_to_rgb_float(hex_color: str):
|
|
hex_color = hex_color.strip().lstrip("#")
|
|
|
|
if len(hex_color) != 6:
|
|
raise ValueError("background_color must be a hex string like #000000")
|
|
|
|
try:
|
|
rgb = tuple(int(hex_color[i : i + 2], 16) / 255.0 for i in (0, 2, 4))
|
|
except ValueError as exc:
|
|
raise ValueError("background_color must be a hex string like #000000") from exc
|
|
|
|
return rgb
|
|
|
|
|
|
class BatchImageGrid:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"columns": (
|
|
"INT",
|
|
{
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"rows": (
|
|
"INT",
|
|
{
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"order": (
|
|
["row_major", "column_major"],
|
|
{
|
|
"default": "row_major",
|
|
},
|
|
),
|
|
"gap": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 256,
|
|
"step": 1,
|
|
},
|
|
),
|
|
"background_color": (
|
|
"STRING",
|
|
{
|
|
"default": "#000000",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "make_grid"
|
|
CATEGORY = "image/grid"
|
|
|
|
def make_grid(self, images, columns, rows, order, gap, background_color):
|
|
if len(images.shape) != 4:
|
|
raise ValueError("images must be an IMAGE tensor with shape [B, H, W, C]")
|
|
|
|
batch_size, height, width, channels = images.shape
|
|
columns = int(columns)
|
|
rows = int(rows)
|
|
gap = int(gap)
|
|
required_count = rows * columns
|
|
|
|
if batch_size < required_count:
|
|
raise ValueError(
|
|
f"Batch Image Grid requires {required_count} images, but received {batch_size}."
|
|
)
|
|
|
|
output_width = (columns * width) + (gap * (columns - 1))
|
|
output_height = (rows * height) + (gap * (rows - 1))
|
|
output = images.new_zeros((1, output_height, output_width, channels))
|
|
|
|
if channels >= 3:
|
|
background = torch.tensor(
|
|
hex_to_rgb_float(background_color),
|
|
dtype=images.dtype,
|
|
device=images.device,
|
|
)
|
|
output[..., :3] = background.view(1, 1, 1, 3)
|
|
else:
|
|
hex_to_rgb_float(background_color)
|
|
|
|
for row in range(rows):
|
|
for col in range(columns):
|
|
if order == "column_major":
|
|
index = (col * rows) + row
|
|
else:
|
|
index = (row * columns) + col
|
|
|
|
y = row * (height + gap)
|
|
x = col * (width + gap)
|
|
output[0, y : y + height, x : x + width, :] = images[index]
|
|
|
|
return (output,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StatusBarOverlay": StatusBarOverlay,
|
|
"KSamplerMultiRun": KSamplerMultiRun,
|
|
"BatchStatusBarOverlay": BatchStatusBarOverlay,
|
|
"BatchImageGrid": BatchImageGrid,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StatusBarOverlay": "Status Bar Overlay",
|
|
"KSamplerMultiRun": "KSampler Multi Run",
|
|
"BatchStatusBarOverlay": "Batch Status Bar Overlay",
|
|
"BatchImageGrid": "Batch Image Grid",
|
|
}
|