feat: sync latest local version as authoritative codebase
Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
This commit is contained in:
o1key
@@ -0,0 +1,241 @@
|
||||
"""
|
||||
高级图像拼接节点
|
||||
支持最多 10 张图像按指定方向（上、下、左、右）依次拼接，
|
||||
支持调整图像大小匹配和添加间隔。
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.file_utils import load_images_from_folder
|
||||
|
||||
|
||||
# 间隔颜色映射
|
||||
SPACING_COLOR_MAP = {
|
||||
"white": (255, 255, 255),
|
||||
"black": (0, 0, 0),
|
||||
"red": (255, 0, 0),
|
||||
"green": (0, 255, 0),
|
||||
"blue": (0, 0, 255),
|
||||
}
|
||||
|
||||
|
||||
def _resize_to_match(img: Image.Image, ref: Image.Image, direction: str) -> Image.Image:
|
||||
"""
|
||||
按拼接方向将 img 缩放，使其与 ref 在垂直于拼接轴的尺寸上一致。
|
||||
|
||||
- 水平拼接 (right/left)：统一高度
|
||||
- 垂直拼接 (down/up)：统一宽度
|
||||
"""
|
||||
ref_w, ref_h = ref.size
|
||||
img_w, img_h = img.size
|
||||
|
||||
if direction in ("right", "left"):
|
||||
if img_h != ref_h:
|
||||
scale = ref_h / img_h
|
||||
new_w = max(1, int(img_w * scale))
|
||||
img = img.resize((new_w, ref_h), Image.LANCZOS)
|
||||
else:
|
||||
if img_w != ref_w:
|
||||
scale = ref_w / img_w
|
||||
new_h = max(1, int(img_h * scale))
|
||||
img = img.resize((ref_w, new_h), Image.LANCZOS)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def _make_spacer(ref: Image.Image, spacing_width: int,
|
||||
direction: str, color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""创建间隔色块"""
|
||||
if direction in ("right", "left"):
|
||||
return Image.new("RGB", (spacing_width, ref.size[1]), color)
|
||||
else:
|
||||
return Image.new("RGB", (ref.size[0], spacing_width), color)
|
||||
|
||||
|
||||
def _stitch_two(img_a: Image.Image, img_b: Image.Image,
|
||||
direction: str, match_size: bool,
|
||||
spacing_width: int, spacing_color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""
|
||||
将两张 PIL 图像按指定方向拼接。
|
||||
img_a 为基准图像，img_b 拼接在 img_a 的指定方向侧。
|
||||
direction="right" → img_b 在 img_a 右侧
|
||||
direction="left" → img_b 在 img_a 左侧
|
||||
direction="down" → img_b 在 img_a 下方
|
||||
direction="up" → img_b 在 img_a 上方
|
||||
"""
|
||||
if img_a.mode != "RGB":
|
||||
img_a = img_a.convert("RGB")
|
||||
if img_b.mode != "RGB":
|
||||
img_b = img_b.convert("RGB")
|
||||
|
||||
if match_size:
|
||||
img_b = _resize_to_match(img_b, img_a, direction)
|
||||
|
||||
if direction == "right":
|
||||
pieces = [img_a, img_b]
|
||||
elif direction == "left":
|
||||
pieces = [img_b, img_a]
|
||||
elif direction == "down":
|
||||
pieces = [img_a, img_b]
|
||||
else: # up
|
||||
pieces = [img_b, img_a]
|
||||
|
||||
if spacing_width > 0:
|
||||
interleaved: List[Image.Image] = []
|
||||
for idx, piece in enumerate(pieces):
|
||||
interleaved.append(piece)
|
||||
if idx < len(pieces) - 1:
|
||||
interleaved.append(_make_spacer(piece, spacing_width, direction, spacing_color))
|
||||
pieces = interleaved
|
||||
|
||||
if direction in ("right", "left"):
|
||||
total_w = sum(p.size[0] for p in pieces)
|
||||
max_h = max(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (total_w, max_h), spacing_color)
|
||||
x = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (x, 0))
|
||||
x += piece.size[0]
|
||||
else:
|
||||
max_w = max(p.size[0] for p in pieces)
|
||||
total_h = sum(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (max_w, total_h), spacing_color)
|
||||
y = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (0, y))
|
||||
y += piece.size[1]
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def _natural_sort_key(filename: str):
|
||||
"""按数字优先的文件名排序，使 1, 2, 3, 10 而非 1, 10, 2, 3"""
|
||||
try:
|
||||
return (0, int(filename))
|
||||
except ValueError:
|
||||
return (1, filename.lower())
|
||||
|
||||
|
||||
class ImageStitchPro:
|
||||
"""
|
||||
高级图像拼接节点
|
||||
|
||||
在 ComfyUI 原生拼接节点基础上扩展，支持同时输入最多 10 张图像，
|
||||
按指定方向依次拼接，并可在图像间添加任意颜色的间隔。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"方向": (["right", "down", "left", "up"], {"default": "down"}),
|
||||
"匹配图像尺寸": ("BOOLEAN", {"default": True}),
|
||||
"间距宽度": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 2}),
|
||||
"间距颜色": (["white", "black", "red", "green", "blue"], {"default": "white"}),
|
||||
},
|
||||
"optional": {
|
||||
"图1": ("IMAGE",),
|
||||
"图2": ("IMAGE",),
|
||||
"图3": ("IMAGE",),
|
||||
"图4": ("IMAGE",),
|
||||
"图5": ("IMAGE",),
|
||||
"图6": ("IMAGE",),
|
||||
"图7": ("IMAGE",),
|
||||
"图8": ("IMAGE",),
|
||||
"图9": ("IMAGE",),
|
||||
"图10": ("IMAGE",),
|
||||
"图11": ("IMAGE",),
|
||||
"图12": ("IMAGE",),
|
||||
"图片路径（可选）": ("STRING", {"default": "", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("拼接图像",)
|
||||
FUNCTION = "stitch"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"高级图像拼接节点，支持最多 12 张图像按指定方向（右/下/左/上）依次拼接。\n"
|
||||
"可选择是否将后续图像缩放以匹配第一张图像的尺寸，并可在图像间添加彩色间隔。\n"
|
||||
"可选填「图片路径」：仅处理该文件夹内图片，按文件名顺序依次拼接；与输入端图片不可同时使用。"
|
||||
)
|
||||
|
||||
def stitch(
|
||||
self,
|
||||
方向: str = "down",
|
||||
匹配图像尺寸: bool = True,
|
||||
间距宽度: int = 0,
|
||||
间距颜色: str = "white",
|
||||
图1: Optional[torch.Tensor] = None,
|
||||
图2: Optional[torch.Tensor] = None,
|
||||
图3: Optional[torch.Tensor] = None,
|
||||
图4: Optional[torch.Tensor] = None,
|
||||
图5: Optional[torch.Tensor] = None,
|
||||
图6: Optional[torch.Tensor] = None,
|
||||
图7: Optional[torch.Tensor] = None,
|
||||
图8: Optional[torch.Tensor] = None,
|
||||
图9: Optional[torch.Tensor] = None,
|
||||
图10: Optional[torch.Tensor] = None,
|
||||
图11: Optional[torch.Tensor] = None,
|
||||
图12: Optional[torch.Tensor] = None,
|
||||
**kwargs: object,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
|
||||
color = SPACING_COLOR_MAP.get(间距颜色, (255, 255, 255))
|
||||
raw_tensors = [图1, 图2, 图3, 图4, 图5, 图6, 图7, 图8, 图9, 图10, 图11, 图12]
|
||||
tensors = [t for t in raw_tensors if t is not None]
|
||||
has_input_images = len(tensors) > 0
|
||||
image_folder = (kwargs.get("图片路径（可选）") or "").strip()
|
||||
|
||||
if image_folder and has_input_images:
|
||||
raise ValueError("不可同时使用「图片路径（可选）」与输入端图片，请二选一。")
|
||||
|
||||
if image_folder:
|
||||
infos = load_images_from_folder(image_folder)
|
||||
if not infos:
|
||||
raise ValueError(f"文件夹中未找到可用的图片，或路径无效: {image_folder}")
|
||||
infos.sort(key=lambda x: _natural_sort_key(x.filename))
|
||||
pil_list = [info.image for info in infos]
|
||||
if len(pil_list) == 1:
|
||||
return (pil_to_tensor(pil_list),)
|
||||
base = pil_list[0]
|
||||
for next_img in pil_list[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
return (pil_to_tensor([base]),)
|
||||
else:
|
||||
if not has_input_images:
|
||||
raise ValueError("请至少接入一张图片，或填写「图片路径（可选）」中的文件夹路径。")
|
||||
|
||||
if len(tensors) == 1:
|
||||
return (tensors[0],)
|
||||
|
||||
pil_batches: List[List[Image.Image]] = [tensor_to_pil(t) for t in tensors]
|
||||
|
||||
batch_size = min(len(b) for b in pil_batches)
|
||||
result_images: List[Image.Image] = []
|
||||
|
||||
for i in range(batch_size):
|
||||
frames = [batch[i] for batch in pil_batches]
|
||||
base = frames[0]
|
||||
for next_img in frames[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
result_images.append(base)
|
||||
|
||||
return (pil_to_tensor(result_images),)
|
||||
Reference in New Issue
Block a user