From 60575a629f9ff31e8fdbf92d44b083211eda327b Mon Sep 17 00:00:00 2001 From: Cyber Dick Lang <286878701@qq.com> Date: Fri, 13 Jun 2025 18:05:17 +0800 Subject: [PATCH] =?UTF-8?q?=E9=9F=B3=E9=A2=91=E8=8A=82=E7=82=B9=E9=87=8D?= =?UTF-8?q?=E6=9E=84=E4=B8=8E=E4=BC=98=E5=8C=96=EF=BC=9A=E5=88=A0=E9=99=A4?= =?UTF-8?q?Upload=E8=8A=82=E7=82=B9=EF=BC=8C=E5=AE=8C=E5=96=84=E4=BE=9D?= =?UTF-8?q?=E8=B5=96=E4=B8=8E=E4=BD=9C=E8=80=85=E4=BF=A1=E6=81=AF=EF=BC=8C?= =?UTF-8?q?=E7=89=88=E6=9C=AC=E6=8F=90=E5=8D=87=E8=87=B31.0.2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- __init__.py | 95 +++++++-- nodes/audio_nodes_utk.py | 136 ++++++++++++ nodes/image_nodes_utk.py | 448 ++++++++++++++++++++++++++++++--------- nodes/tool_nodes_utk.py | 28 +++ requirements.txt | 6 +- 5 files changed, 590 insertions(+), 123 deletions(-) create mode 100644 nodes/audio_nodes_utk.py diff --git a/__init__.py b/__init__.py index 720f079..1ef88d3 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,49 @@ -from .nodes import ( - EmptyUnitGenerator_UTK, - ImageRatioDetector_UTK, - ShowInt_UTK, - ShowFloat_UTK, - ShowList_UTK, - ShowText_UTK, - PreviewMask_UTK, - DepthMapBlur_UTK, - ImageConcatenate_UTK, - ImageConcatenateMulti_UTK, -) +""" +ComfyUI Universal Toolkit +~~~~~~~~~~~~~~~~~~~~~~~ + +A comprehensive toolkit for ComfyUI that provides various utility nodes for image processing, text manipulation, and more. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +__version__ = "1.0.2" +__author__ = "CyberDickLang" +__email__ = "286878701@qq.com" +__url__ = "https://github.com/whmc76" + +# 更新日志 +CHANGELOG = { + "1.0.2": [ + "删除无用节点 Load Audio Plus Upload (UTK)", + "新增 Audio Crop Process (UTK) 支持原生音频上传", + "修复音频处理相关bug，完善依赖" + ], + "1.0.1": [ + "改进 ImageConcatenate_UTK 节点：", + "- 添加 match_image_size 参数，支持自动匹配图像尺寸", + "- 添加 max_size 参数，限制最大输出尺寸", + "- 添加背景颜色选项", + "- 优化图像拼接逻辑，保持宽高比", + "改进 ImageConcatenateMulti_UTK 节点：", + "- 与 ImageConcatenate_UTK 保持一致的功能", + "- 支持自动方向选择", + "- 支持图像尺寸匹配", + "- 支持最大尺寸限制", + "- 优化多图拼接逻辑" + ], + "1.0.0": [ + "初始版本发布", + "包含基础图像处理节点", + "包含文本处理节点", + "包含工具类节点" + ] +} + +from .nodes.image_nodes_utk import EmptyUnitGenerator_UTK, ImageRatioDetector_UTK, DepthMapBlur_UTK, ImageConcatenate_UTK, ImageConcatenateMulti_UTK +from .nodes.tool_nodes_utk import ShowInt_UTK, ShowFloat_UTK, ShowList_UTK, ShowText_UTK, PreviewMask_UTK +from .nodes.audio_nodes_utk import LoadAudioPlusFromPath_UTK, AudioCropProcessUTK NODE_CLASS_MAPPINGS = { "EmptyUnitGenerator_UTK": EmptyUnitGenerator_UTK, @@ -22,19 +56,44 @@ NODE_CLASS_MAPPINGS = { "DepthMapBlur_UTK": DepthMapBlur_UTK, "ImageConcatenate_UTK": ImageConcatenate_UTK, "ImageConcatenateMulti_UTK": ImageConcatenateMulti_UTK, + "LoadAudioPlusFromPath_UTK": LoadAudioPlusFromPath_UTK, + "AudioCropProcessUTK": AudioCropProcessUTK, } NODE_DISPLAY_NAME_MAPPINGS = { - "EmptyUnitGenerator_UTK": "Empty Unit Generator (UTK)", - "ImageRatioDetector_UTK": "Image Ratio Detector (UTK)", + "EmptyUnitGenerator_UTK": "Empty Unit Generator", + "ImageRatioDetector_UTK": "Image Ratio Detector", "ShowInt_UTK": "Show Int (UTK)", "ShowFloat_UTK": "Show Float (UTK)", "ShowList_UTK": "Show List (UTK)", "ShowText_UTK": "Show Text (UTK)", "PreviewMask_UTK": "Preview Mask (UTK)", - "DepthMapBlur_UTK": "Depth Map Blur (UTK)", - "ImageConcatenate_UTK": "Image Concatenate (UTK)", - "ImageConcatenateMulti_UTK": "Image Concatenate Multi (UTK)", + "DepthMapBlur_UTK": "Depth Map Blur", + "ImageConcatenate_UTK": "Image Concatenate", + "ImageConcatenateMulti_UTK": "Image Concatenate Multi", + "LoadAudioPlusFromPath_UTK": "Load Audio Plus From Path (UTK)", + "AudioCropProcessUTK": "Audio Crop Process (UTK)", } -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +NODE_CATEGORIES = { + "UniversalToolkit": [ + "EmptyUnitGenerator_UTK", + "ImageRatioDetector_UTK", + "DepthMapBlur_UTK", + "ImageConcatenate_UTK", + "ImageConcatenateMulti_UTK", + "LoadAudioPlusFromPath_UTK", + "AudioCropProcessUTK", + ] +} + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "NODE_CATEGORIES", + "__version__", + "__author__", + "__email__", + "__url__", + "CHANGELOG" +] \ No newline at end of file diff --git a/nodes/audio_nodes_utk.py b/nodes/audio_nodes_utk.py new file mode 100644 index 0000000..3f04e3d --- /dev/null +++ b/nodes/audio_nodes_utk.py @@ -0,0 +1,136 @@ +import math +import os +from pathlib import Path +import torch +import soundfile as sf +import librosa + +FLOAT_MAX = 99999999999999999.0 + +class LoadAudioPlusFromPath_UTK: + CATEGORY = "UniversalToolkit" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "path": ("STRING", {"default": "./audio.mp3"}), + "gain_db": ("FLOAT", {"default": 0, "min": -100, "max": 100}), + "offset_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "duration_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "resample_to_hz": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "make_stereo": ("BOOLEAN", {"default": True}), + } + } + RETURN_TYPES = ("AUDIO", "INT", "INT", "FLOAT") + RETURN_NAMES = ("audio", "sample_rate", "channels", "duration") + FUNCTION = "execute" + + def db_to_scalar(self, db: float): + return 10 ** (db / 20) + + @classmethod + def IS_CHANGED(cls, path: str, *args): + if os.path.exists(path): + mtime = os.path.getmtime(path) + else: + mtime = None + return (mtime, path, *args) + + def execute( + self, + path: str, + gain_db: float, + offset_seconds: float, + duration_seconds: float, + resample_to_hz: float, + make_stereo: bool, + ): + # 路径预处理：去除首尾单双引号，替换分隔符，兼容Windows绝对路径 + path = path.strip().strip('"').strip("'") + path = path.replace('\\', '/').replace('\\', '/') + if os.name == 'nt' and len(path) > 2 and path[1] == ':': + path = path.replace('\\', '/') + # 文件存在性检查，异常时详细提示 + if not os.path.isfile(path): + raise FileNotFoundError(f"音频文件不存在或路径错误: {path}\n请检查路径是否正确，注意不要包含多余的引号或空格，Windows下建议使用/或\\分隔符。") + # 加载音频，异常时详细提示 + try: + sr = int(resample_to_hz) if resample_to_hz > 0 else None + duration = duration_seconds if duration_seconds > 0 else None + mix, sr = librosa.load(path, sr=sr, mono=False, offset=offset_seconds, duration=duration) + except Exception as e: + raise RuntimeError(f"音频加载失败: {e}\n请确认文件格式是否受支持，路径是否包含特殊字符。") + # shape调整 + if len(mix.shape) == 1: + mix = torch.stack([mix], dim=0) + if make_stereo: + if mix.shape[0] == 1: + mix = torch.cat([mix, mix], dim=0) + elif mix.shape[0] == 2: + pass + else: + raise ValueError(f"Input audio has {mix.shape[0]} channels, cannot convert to stereo (2 channels)") + mix = torch.from_numpy(mix) # 保证为Tensor + mix = torch.unsqueeze(mix, 0) # shape: [1, 2, N] 或 [1, 1, N] + if gain_db != 0.0: + gain_scalar = 10 ** (gain_db / 20) + mix = gain_scalar * mix + sample_rate = int(sr) + channels = int(mix.shape[1]) + duration_val = float(mix.shape[2] / sample_rate) if sample_rate > 0 else 0.0 + return ({"sample_rate": sample_rate, "waveform": mix}, sample_rate, channels, duration_val) + +class AudioCropProcessUTK: + CATEGORY = "UniversalToolkit" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "gain_db": ("FLOAT", {"default": 0, "min": -100, "max": 100}), + "offset_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "duration_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "resample_to_hz": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}), + "make_stereo": ("BOOLEAN", {"default": True}), + } + } + RETURN_TYPES = ("AUDIO", "INT", "INT", "FLOAT") + RETURN_NAMES = ("audio", "sample_rate", "channels", "duration") + FUNCTION = "execute" + + @classmethod + def IS_CHANGED(cls, *args): + return args + + def execute( + self, + audio, + gain_db: float, + offset_seconds: float, + duration_seconds: float, + resample_to_hz: float, + make_stereo: bool, + ): + waveform = audio["waveform"] # [B, C, N] + sample_rate = int(audio["sample_rate"]) + # 裁剪offset和duration + start = int(offset_seconds * sample_rate) + end = int(start + duration_seconds * sample_rate) if duration_seconds > 0 else waveform.shape[2] + waveform = waveform[:, :, start:end] + # 重采样 + if resample_to_hz > 0 and int(resample_to_hz) != sample_rate: + import torchaudio + waveform = torchaudio.functional.resample(waveform, sample_rate, int(resample_to_hz)) + sample_rate = int(resample_to_hz) + # 增益 + if gain_db != 0.0: + gain_scalar = 10 ** (gain_db / 20) + waveform = waveform * gain_scalar + # 强制立体声 + if make_stereo and waveform.shape[1] == 1: + waveform = torch.cat([waveform, waveform], dim=1) + elif make_stereo and waveform.shape[1] != 2: + raise ValueError(f"Input audio has {waveform.shape[1]} channels, cannot convert to stereo (2 channels)") + channels = int(waveform.shape[1]) + duration_val = float(waveform.shape[2] / sample_rate) if sample_rate > 0 else 0.0 + return ({"sample_rate": sample_rate, "waveform": waveform}, sample_rate, channels, duration_val) \ No newline at end of file diff --git a/nodes/image_nodes_utk.py b/nodes/image_nodes_utk.py index 7f66ec7..b0eee4a 100644 --- a/nodes/image_nodes_utk.py +++ b/nodes/image_nodes_utk.py @@ -2,6 +2,8 @@ import torch import numpy as np import re import math +from comfy.utils import ProgressBar, common_upscale +from PIL import Image class EmptyUnitGenerator_UTK: CATEGORY = "UniversalToolkit" @@ -225,82 +227,200 @@ class ImageConcatenate_UTK: "required": { "image1": ("IMAGE",), "image2": ("IMAGE",), - "direction": (["right", "left", "up", "down", "auto"], {"default": "auto"}), - "align": (["start", "center", "end"], {"default": "center"}), - "spacing": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 1}), + "direction": ( + [ 'right', + 'down', + 'left', + 'up', + 'auto', + ], + { + "default": 'auto' + }), + "match_image_size": ("BOOLEAN", {"default": True}), + "max_size": ("INT", {"default": 4096, "min": 64, "max": 8192, "step": 64}), + "background_color": (["black", "white", "gray", "transparent"], {"default": "black"}), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "concatenate" - def concatenate(self, image1, image2, direction, align, spacing): - # 获取图像尺寸 - b1, c1, h1, w1 = image1.shape - b2, c2, h2, w2 = image2.shape - - # 确保通道数相同 - if c1 != c2: - raise ValueError("输入图像的通道数必须相同") - - # 如果是自动模式，根据图像尺寸决定拼接方向 - if direction == "auto": - # 计算水平拼接和垂直拼接后的宽高比 - horizontal_ratio = (w1 + w2 + spacing) / max(h1, h2) - vertical_ratio = max(w1, w2) / (h1 + h2 + spacing) + def concatenate(self, image1, image2, direction, match_image_size, max_size, background_color): + # Check if the batch sizes are different + batch_size1 = image1.shape[0] + batch_size2 = image2.shape[0] + + if batch_size1 != batch_size2: + # Calculate the number of repetitions needed + max_batch_size = max(batch_size1, batch_size2) + repeats1 = max_batch_size - batch_size1 + repeats2 = max_batch_size - batch_size2 - # 选择更接近1:1比例的拼接方式 - direction = "right" if abs(horizontal_ratio - 1) <= abs(vertical_ratio - 1) else "down" - - # 计算输出尺寸 - if direction in ["right", "left"]: - out_h = max(h1, h2) - out_w = w1 + w2 + spacing + # Repeat the last image to match the largest batch size + if repeats1 > 0: + last_image1 = image1[-1].unsqueeze(0).repeat(repeats1, 1, 1, 1) + image1 = torch.cat([image1.clone(), last_image1], dim=0) + if repeats2 > 0: + last_image2 = image2[-1].unsqueeze(0).repeat(repeats2, 1, 1, 1) + image2 = torch.cat([image2.clone(), last_image2], dim=0) + + # Get original dimensions + h1, w1 = image1.shape[1:3] + h2, w2 = image2.shape[1:3] + + # If direction is auto, determine the best direction based on image dimensions + if direction == 'auto': + # Calculate aspect ratios for both directions + horizontal_ratio = (w1 + w2) / max(h1, h2) + vertical_ratio = max(w1, w2) / (h1 + h2) + + # Choose the direction that results in an aspect ratio closer to 1:1 + direction = 'right' if abs(horizontal_ratio - 1) <= abs(vertical_ratio - 1) else 'down' + + # Match image sizes if requested + if match_image_size: + if direction in ['right', 'left', 'auto']: + # Match heights + target_height = max(h1, h2) + # Scale image1 if needed + if h1 < target_height: + scale = target_height / h1 + new_width = int(w1 * scale) + image1 = torch.nn.functional.interpolate( + image1.permute(0, 3, 1, 2), + size=(target_height, new_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + # Scale image2 if needed + if h2 < target_height: + scale = target_height / h2 + new_width = int(w2 * scale) + image2 = torch.nn.functional.interpolate( + image2.permute(0, 3, 1, 2), + size=(target_height, new_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + else: # up, down + # Match widths + target_width = max(w1, w2) + # Scale image1 if needed + if w1 < target_width: + scale = target_width / w1 + new_height = int(h1 * scale) + image1 = torch.nn.functional.interpolate( + image1.permute(0, 3, 1, 2), + size=(new_height, target_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + # Scale image2 if needed + if w2 < target_width: + scale = target_width / w2 + new_height = int(h2 * scale) + image2 = torch.nn.functional.interpolate( + image2.permute(0, 3, 1, 2), + size=(new_height, target_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + + # Update dimensions after scaling + h1, w1 = image1.shape[1:3] + h2, w2 = image2.shape[1:3] + + # Calculate final dimensions + if direction in ['right', 'left']: + final_height = max(h1, h2) + final_width = w1 + w2 else: # up, down - out_h = h1 + h2 + spacing - out_w = max(w1, w2) - - # 创建输出张量 - output = torch.zeros((b1, c1, out_h, out_w), dtype=image1.dtype, device=image1.device) - - # 计算对齐位置 - if direction in ["right", "left"]: - if align == "start": - y1 = y2 = 0 - elif align == "center": - y1 = (out_h - h1) // 2 - y2 = (out_h - h2) // 2 - else: # end - y1 = out_h - h1 - y2 = out_h - h2 + final_height = h1 + h2 + final_width = max(w1, w2) + + # Check if we need to scale down + if max(final_height, final_width) > max_size: + scale = max_size / max(final_height, final_width) + new_h1 = int(h1 * scale) + new_w1 = int(w1 * scale) + new_h2 = int(h2 * scale) + new_w2 = int(w2 * scale) - if direction == "right": - x1 = 0 - x2 = w1 + spacing - else: # left - x1 = w2 + spacing - x2 = 0 - else: # up, down - if align == "start": - x1 = x2 = 0 - elif align == "center": - x1 = (out_w - w1) // 2 - x2 = (out_w - w2) // 2 - else: # end - x1 = out_w - w1 - x2 = out_w - w2 + # Scale both images + image1 = torch.nn.functional.interpolate( + image1.permute(0, 3, 1, 2), + size=(new_h1, new_w1), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) - if direction == "down": - y1 = 0 - y2 = h1 + spacing - else: # up - y1 = h2 + spacing - y2 = 0 - - # 复制图像数据 - output[:, :, y1:y1+h1, x1:x1+w1] = image1 - output[:, :, y2:y2+h2, x2:x2+w2] = image2 - + image2 = torch.nn.functional.interpolate( + image2.permute(0, 3, 1, 2), + size=(new_h2, new_w2), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + + # Update dimensions after scaling + h1, w1 = image1.shape[1:3] + h2, w2 = image2.shape[1:3] + + # Recalculate final dimensions + if direction in ['right', 'left']: + final_height = max(h1, h2) + final_width = w1 + w2 + else: # up, down + final_height = h1 + h2 + final_width = max(w1, w2) + + # Ensure both images have the same number of channels + channels_image1 = image1.shape[-1] + channels_image2 = image2.shape[-1] + + if channels_image1 != channels_image2: + if channels_image1 < channels_image2: + # Add alpha channel to image1 if image2 has it + alpha_channel = torch.ones((*image1.shape[:-1], channels_image2 - channels_image1), device=image1.device) + image1 = torch.cat((image1, alpha_channel), dim=-1) + else: + # Add alpha channel to image2 if image1 has it + alpha_channel = torch.ones((*image2.shape[:-1], channels_image1 - channels_image2), device=image2.device) + image2 = torch.cat((image2, alpha_channel), dim=-1) + + # Create output tensor with background color + if background_color == "transparent": + output = torch.zeros((batch_size1, final_height, final_width, channels_image1), dtype=image1.dtype, device=image1.device) + else: + color_value = 1.0 if background_color == "white" else 0.0 if background_color == "black" else 0.5 + output = torch.full((batch_size1, final_height, final_width, channels_image1), color_value, dtype=image1.dtype, device=image1.device) + + # Calculate positions for image placement + if direction == 'right': + x1 = 0 + x2 = w1 + y1 = (final_height - h1) // 2 + y2 = (final_height - h2) // 2 + elif direction == 'left': + x1 = w2 + x2 = 0 + y1 = (final_height - h1) // 2 + y2 = (final_height - h2) // 2 + elif direction == 'down': + x1 = (final_width - w1) // 2 + x2 = (final_width - w2) // 2 + y1 = 0 + y2 = h1 + else: # up + x1 = (final_width - w1) // 2 + x2 = (final_width - w2) // 2 + y1 = h2 + y2 = 0 + + # Place images in the output tensor + output[:, y1:y1+h1, x1:x1+w1] = image1 + output[:, y2:y2+h2, x2:x2+w2] = image2 + return (output,) class ImageConcatenateMulti_UTK: @@ -311,51 +431,173 @@ class ImageConcatenateMulti_UTK: return { "required": { "images": ("IMAGE",), - "direction": (["horizontal", "vertical"], {"default": "horizontal"}), - "align": (["start", "center", "end"], {"default": "center"}), - "spacing": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 1}), + "direction": ( + [ 'right', + 'down', + 'left', + 'up', + 'auto', + ], + { + "default": 'auto' + }), + "match_image_size": ("BOOLEAN", {"default": True}), + "max_size": ("INT", {"default": 4096, "min": 64, "max": 8192, "step": 64}), + "background_color": (["black", "white", "gray", "transparent"], {"default": "black"}), + "grid_size": (["auto", "1x1", "2x2", "3x3", "4x4"], {"default": "auto"}), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "concatenate_multi" - def concatenate_multi(self, images, direction, align, spacing): + def concatenate_multi(self, images, direction, match_image_size, max_size, background_color, grid_size): if len(images.shape) != 4: - raise ValueError("输入必须是4D张量 [batch, channels, height, width]") + raise ValueError("输入必须是4D张量 [batch, height, width, channels]") - b, c, h, w = images.shape - - # 计算输出尺寸 - if direction == "horizontal": - out_h = h - out_w = w * b + spacing * (b - 1) - else: # vertical - out_h = h * b + spacing * (b - 1) - out_w = w + batch_size = images.shape[0] + # 处理网格布局 + if grid_size != "auto": + rows, cols = map(int, grid_size.split("x")) + if batch_size > rows * cols: + raise ValueError(f"图像数量 ({batch_size}) 超过网格大小 ({grid_size})") + # 填充到网格大小 + if batch_size < rows * cols: + padding = torch.zeros((rows * cols - batch_size, *images.shape[1:]), dtype=images.dtype, device=images.device) + images = torch.cat([images, padding], dim=0) + batch_size = rows * cols + else: + # 自动计算网格大小 + if direction in ['right', 'left', 'auto']: + rows = 1 + cols = batch_size + else: # up, down + rows = batch_size + cols = 1 + + # 获取所有图像的尺寸 + heights = [] + widths = [] + for i in range(batch_size): + h, w = images[i].shape[:2] + heights.append(h) + widths.append(w) + + # 如果方向是auto，计算最佳方向 + if direction == 'auto': + # 计算水平和垂直拼接的宽高比 + total_width = sum(widths) + max_height = max(heights) + horizontal_ratio = total_width / max_height + + total_height = sum(heights) + max_width = max(widths) + vertical_ratio = max_width / total_height + + # 选择更接近1:1的方向 + direction = 'right' if abs(horizontal_ratio - 1) <= abs(vertical_ratio - 1) else 'down' + + # 如果需要匹配图像尺寸 + if match_image_size: + if direction in ['right', 'left', 'auto']: + # 匹配高度 + target_height = max(heights) + for i in range(batch_size): + if heights[i] < target_height: + scale = target_height / heights[i] + new_width = int(widths[i] * scale) + images[i] = torch.nn.functional.interpolate( + images[i].unsqueeze(0).permute(0, 3, 1, 2), + size=(target_height, new_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1).squeeze(0) + else: # up, down + # 匹配宽度 + target_width = max(widths) + for i in range(batch_size): + if widths[i] < target_width: + scale = target_width / widths[i] + new_height = int(heights[i] * scale) + images[i] = torch.nn.functional.interpolate( + images[i].unsqueeze(0).permute(0, 3, 1, 2), + size=(new_height, target_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1).squeeze(0) + + # 更新尺寸 + heights = [] + widths = [] + for i in range(batch_size): + h, w = images[i].shape[:2] + heights.append(h) + widths.append(w) + + # 计算最终输出尺寸 + if direction in ['right', 'left', 'auto']: + final_height = max(heights) + final_width = sum(widths) + else: # up, down + final_height = sum(heights) + final_width = max(widths) + + # 检查是否需要缩放 + if max(final_height, final_width) > max_size: + scale = max_size / max(final_height, final_width) + for i in range(batch_size): + new_height = int(heights[i] * scale) + new_width = int(widths[i] * scale) + images[i] = torch.nn.functional.interpolate( + images[i].unsqueeze(0).permute(0, 3, 1, 2), + size=(new_height, new_width), + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1).squeeze(0) + + # 更新最终尺寸 + heights = [] + widths = [] + for i in range(batch_size): + h, w = images[i].shape[:2] + heights.append(h) + widths.append(w) + + if direction in ['right', 'left', 'auto']: + final_height = max(heights) + final_width = sum(widths) + else: # up, down + final_height = sum(heights) + final_width = max(widths) + # 创建输出张量 - output = torch.zeros((1, c, out_h, out_w), dtype=images.dtype, device=images.device) - - # 复制图像数据 - for i in range(b): - if direction == "horizontal": - if align == "start": - y = 0 - elif align == "center": - y = (out_h - h) // 2 - else: # end - y = out_h - h - x = i * (w + spacing) - else: # vertical - if align == "start": - x = 0 - elif align == "center": - x = (out_w - w) // 2 - else: # end - x = out_w - w - y = i * (h + spacing) - - output[0, :, y:y+h, x:x+w] = images[i] - + if background_color == "transparent": + output = torch.zeros((1, final_height, final_width, images.shape[-1]), dtype=images.dtype, device=images.device) + else: + color_value = 1.0 if background_color == "white" else 0.0 if background_color == "black" else 0.5 + output = torch.full((1, final_height, final_width, images.shape[-1]), color_value, dtype=images.dtype, device=images.device) + + # 放置图像 + if direction in ['right', 'left', 'auto']: + x_offset = 0 + for i in range(batch_size): + h, w = images[i].shape[:2] + y_offset = (final_height - h) // 2 + if direction == 'left': + x_offset = final_width - sum(widths[i:]) + output[0, y_offset:y_offset+h, x_offset:x_offset+w] = images[i] + if direction != 'left': + x_offset += w + else: # up, down + y_offset = 0 + for i in range(batch_size): + h, w = images[i].shape[:2] + x_offset = (final_width - w) // 2 + if direction == 'up': + y_offset = final_height - sum(heights[i:]) + output[0, y_offset:y_offset+h, x_offset:x_offset+w] = images[i] + if direction != 'up': + y_offset += h + return (output,) \ No newline at end of file diff --git a/nodes/tool_nodes_utk.py b/nodes/tool_nodes_utk.py index 1f03a98..bd1bb28 100644 --- a/nodes/tool_nodes_utk.py +++ b/nodes/tool_nodes_utk.py @@ -1,5 +1,33 @@ import torch +class Show_UTK: + CATEGORY = "UniversalToolkit" + @classmethod + def INPUT_TYPES(cls): + return {"required": {"input": ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")}} + RETURN_TYPES = ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT") + RETURN_NAMES = ("string", "int", "float", "list", "mask", "image", "latent") + FUNCTION = "show" + IS_PREVIEW = True + + def show(self, input): + outs = [None] * 7 + if isinstance(input, str): + outs[0] = input + elif isinstance(input, int): + outs[1] = input + elif isinstance(input, float): + outs[2] = input + elif isinstance(input, list): + outs[3] = input + elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 1: + outs[4] = input # MASK + elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 3: + outs[5] = input # IMAGE + elif isinstance(input, dict) and "samples" in input: + outs[6] = input # LATENT + return tuple(outs) + class ShowInt_UTK: CATEGORY = "UniversalToolkit" @classmethod diff --git a/requirements.txt b/requirements.txt index 8323421..6d9bb35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ Pillow -numpy>=1.24.0 -torch>=2.0.0 \ No newline at end of file +numpy +torch +librosa +torchaudio \ No newline at end of file