Hawk Lee
e21a5c9eba
The eff537d commit changed all files from 755 to 644.
Restores original executable permissions.
257 lines
14 KiB
Python
Executable File
257 lines
14 KiB
Python
Executable File
import torch
|
||
import os
|
||
import shutil # 用于清理临时目录 (如果需要)
|
||
import time
|
||
from PIL import Image, ImageOps # Pillow 用于图像处理
|
||
import numpy as np
|
||
import folder_paths # 用于获取输出目录
|
||
from comfy.utils import ProgressBar
|
||
from tqdm import tqdm
|
||
import glob # 用于查找文件
|
||
|
||
# 可能需要 common_upscale 类似的缩放，但我们会用 Pillow 实现
|
||
# from comfy.model_management import common_upscale # 如果你打算复用它，但对于磁盘文件，Pillow更直接
|
||
|
||
class AIIA_Utils_Image_Concanate:
|
||
NODE_NAME = "AIIA Utils Image Concatenate (Disk)" # 加上 (Disk) 以示区别
|
||
CATEGORY = "AIIA/Utils" # 或者 AIIA/Image Utils
|
||
FUNCTION = "concatenate_images_from_disk"
|
||
RETURN_TYPES = ("STRING", "INT") # output_directory, frame_count
|
||
RETURN_NAMES = ("output_directory", "concatenated_frame_count")
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"directory_1": ("STRING", {"default": "path/to/frames1", "multiline": False}),
|
||
"directory_2": ("STRING", {"default": "path/to/frames2", "multiline": False}),
|
||
"direction": (['right', 'down', 'left', 'up'], {"default": 'right'}),
|
||
"match_image_size": ("BOOLEAN", {"default": True}),
|
||
# 如果 match_image_size 为 True，以哪个目录的图像尺寸为基准
|
||
"base_size_from": (['directory_1', 'directory_2'], {"default": 'directory_1'}),
|
||
"output_subdir_name": ("STRING", {"default": "concatenated_frames_AIIA"}),
|
||
},
|
||
"optional": {
|
||
# 可以添加填充颜色等
|
||
"background_color": ("STRING", {"default": "white", "tooltip": "背景颜色 (例如 'white', 'black', '#RRGGBB')"}),
|
||
}
|
||
}
|
||
|
||
def _get_sorted_image_files(self, directory_path):
|
||
"""辅助函数：获取目录下排序后的图像文件列表"""
|
||
# 支持常见的图像格式，扩展大小写支持以适应 Linux
|
||
supported_extensions = ["*.png", "*.PNG", "*.jpg", "*.JPG", "*.jpeg", "*.JPEG", "*.bmp", "*.BMP", "*.webp", "*.WEBP"]
|
||
image_files = []
|
||
for ext in supported_extensions:
|
||
image_files.extend(glob.glob(os.path.join(directory_path, ext)))
|
||
|
||
# 按文件名排序 (假设文件名包含数字序列如 frame_001.png)
|
||
# 更稳健的排序可能需要解析文件名中的数字
|
||
try:
|
||
# 尝试基于文件名中的数字进行自然排序
|
||
image_files.sort(key=lambda f: int("".join(filter(str.isdigit, os.path.basename(f))) or 0))
|
||
except ValueError:
|
||
# 如果无法解析数字，则按普通字符串排序
|
||
image_files.sort()
|
||
print(f"警告: 目录 {directory_path} 中的文件无法按数字序列智能排序，将使用标准字符串排序。")
|
||
|
||
if not image_files:
|
||
print(f"警告: 目录 {directory_path} 中未找到支持的图像文件。")
|
||
return image_files
|
||
|
||
def concatenate_images_from_disk(self, directory_1, directory_2, direction,
|
||
match_image_size, base_size_from, output_subdir_name,
|
||
background_color="white"):
|
||
node_name_log = f"[{self.__class__.NODE_NAME}]"
|
||
print(f"{node_name_log} 开始拼接图像 (从磁盘)。")
|
||
start_time_process = time.time()
|
||
|
||
files1 = self._get_sorted_image_files(directory_1)
|
||
files2 = self._get_sorted_image_files(directory_2)
|
||
|
||
num_frames1 = len(files1)
|
||
num_frames2 = len(files2)
|
||
print(f"{node_name_log} DEBUG: 目录1 ('{os.path.basename(directory_1)}') 包含 {num_frames1} 个图像文件。")
|
||
print(f"{node_name_log} DEBUG: 目录2 ('{os.path.basename(directory_2)}') 包含 {num_frames2} 个图像文件。")
|
||
|
||
if not files1 or not files2:
|
||
error_msg = "输入目录为空或未找到图像文件。"
|
||
if not files1: error_msg += f" 检查目录1: {directory_1}"
|
||
if not files2: error_msg += f" 检查目录2: {directory_2}"
|
||
print(f"错误: {node_name_log} {error_msg}")
|
||
return (f"错误: {error_msg}", 0)
|
||
|
||
# --- 处理帧数不匹配 ---
|
||
if num_frames1 != num_frames2:
|
||
# 当前简单处理：按最短的序列处理
|
||
num_frames_to_process = min(num_frames1, num_frames2)
|
||
print(f"警告: {node_name_log} 输入目录帧数不匹配 ({num_frames1} vs {num_frames2})。将处理 {num_frames_to_process} 帧。")
|
||
files1 = files1[:num_frames_to_process]
|
||
files2 = files2[:num_frames_to_process]
|
||
else:
|
||
num_frames_to_process = num_frames1
|
||
|
||
if num_frames_to_process == 0:
|
||
print(f"错误: {node_name_log} 没有可供处理的匹配帧。")
|
||
return ("错误: 没有可供处理的匹配帧。", 0)
|
||
|
||
# --- 创建输出目录 ---
|
||
output_main_dir = folder_paths.get_output_directory()
|
||
timestamp_str = time.strftime("%Y%m%d-%H%M%S")
|
||
unique_folder_name = f"{output_subdir_name}_{timestamp_str}_{int(torch.randint(0,10000,(1,)).item())}"
|
||
final_output_dir = os.path.join(output_main_dir, unique_folder_name)
|
||
try:
|
||
os.makedirs(final_output_dir, exist_ok=True)
|
||
except Exception as e_mkdir:
|
||
print(f"错误: {node_name_log} 无法创建输出目录 {final_output_dir}: {e_mkdir}")
|
||
return (f"错误: 无法创建输出目录: {e_mkdir}", 0)
|
||
|
||
concatenated_count = 0
|
||
comfy_pbar = ProgressBar(num_frames_to_process)
|
||
|
||
print(f"{node_name_log} 将处理 {num_frames_to_process} 帧组。")
|
||
|
||
with tqdm(total=num_frames_to_process, desc=f"{node_name_log} Concatenating", unit="frame") as console_pbar:
|
||
for i in range(num_frames_to_process):
|
||
try:
|
||
img1_pil = Image.open(files1[i])
|
||
img2_pil = Image.open(files2[i])
|
||
|
||
# 确保图像是 RGB 或 RGBA，如果一个是P模式（调色板），先转换
|
||
if img1_pil.mode == 'P': img1_pil = img1_pil.convert('RGBA' if 'A' in img2_pil.mode else 'RGB')
|
||
if img2_pil.mode == 'P': img2_pil = img2_pil.convert('RGBA' if 'A' in img1_pil.mode else 'RGB')
|
||
|
||
# 统一通道（例如，如果一个是RGB，一个是RGBA，都转为RGBA）
|
||
if 'A' in img1_pil.mode and 'A' not in img2_pil.mode:
|
||
img2_pil = img2_pil.convert('RGBA')
|
||
elif 'A' in img2_pil.mode and 'A' not in img1_pil.mode:
|
||
img1_pil = img1_pil.convert('RGBA')
|
||
|
||
# 保留原始图像副本，以备尺寸匹配时参考
|
||
img1_orig_pil = img1_pil.copy()
|
||
img2_orig_pil = img2_pil.copy()
|
||
|
||
|
||
# --- 尺寸匹配逻辑 (借鉴并用Pillow实现) ---
|
||
if match_image_size:
|
||
if base_size_from == 'directory_1':
|
||
base_img_pil = img1_orig_pil
|
||
target_img_pil = img2_pil # 我们要修改 img2_pil
|
||
target_img_orig_pil = img2_orig_pil # 用于计算宽高比
|
||
else: # base_size_from == 'directory_2'
|
||
base_img_pil = img2_orig_pil
|
||
target_img_pil = img1_pil # 我们要修改 img1_pil
|
||
target_img_orig_pil = img1_orig_pil # 用于计算宽高比
|
||
|
||
base_w, base_h = base_img_pil.size
|
||
target_orig_w, target_orig_h = target_img_orig_pil.size
|
||
|
||
if target_orig_h == 0 or target_orig_w == 0: # 避免除以零
|
||
print(f"警告: {node_name_log} 第 {i} 帧中, 目标图像尺寸为零，跳过尺寸匹配。")
|
||
else:
|
||
aspect_ratio = target_orig_w / target_orig_h
|
||
|
||
if direction in ['left', 'right']: # 匹配高度，调整宽度
|
||
new_h = base_h
|
||
new_w = int(new_h * aspect_ratio)
|
||
else: # up, down: 匹配宽度，调整高度
|
||
new_w = base_w
|
||
new_h = int(new_w / aspect_ratio)
|
||
|
||
if new_w <=0 or new_h <=0: # 避免无效尺寸
|
||
print(f"警告: {node_name_log} 第 {i} 帧中, 计算出的新尺寸无效 ({new_w}x{new_h})，跳过此帧的尺寸匹配。")
|
||
else:
|
||
# 使用 LANCZOS (高质量) 或 ANTIALIAS (Pillow 9.0.0+ 推荐)
|
||
resample_filter = Image.Resampling.LANCZOS if hasattr(Image.Resampling, 'LANCZOS') else Image.LANCZOS
|
||
resized_target_img = target_img_orig_pil.resize((new_w, new_h), resample=resample_filter)
|
||
|
||
# 更新 PIL 对象
|
||
if base_size_from == 'directory_1':
|
||
img2_pil = resized_target_img
|
||
else:
|
||
img1_pil = resized_target_img
|
||
|
||
# --- 拼接逻辑 ---
|
||
w1, h1 = img1_pil.size
|
||
w2, h2 = img2_pil.size
|
||
|
||
if direction == 'right':
|
||
new_width = w1 + w2
|
||
new_height = max(h1, h2)
|
||
# 创建新画布，使用用户指定的背景色
|
||
result_img = Image.new(img1_pil.mode, (new_width, new_height), background_color)
|
||
result_img.paste(img1_pil, (0, (new_height - h1) // 2)) # 居中粘贴
|
||
result_img.paste(img2_pil, (w1, (new_height - h2) // 2))
|
||
elif direction == 'down':
|
||
new_width = max(w1, w2)
|
||
new_height = h1 + h2
|
||
result_img = Image.new(img1_pil.mode, (new_width, new_height), background_color)
|
||
result_img.paste(img1_pil, ((new_width - w1) // 2, 0))
|
||
result_img.paste(img2_pil, ((new_width - w2) // 2, h1))
|
||
elif direction == 'left':
|
||
new_width = w1 + w2
|
||
new_height = max(h1, h2)
|
||
result_img = Image.new(img1_pil.mode, (new_width, new_height), background_color)
|
||
result_img.paste(img2_pil, (0, (new_height - h2) // 2))
|
||
result_img.paste(img1_pil, (w2, (new_height - h1) // 2))
|
||
elif direction == 'up':
|
||
new_width = max(w1, w2)
|
||
new_height = h1 + h2
|
||
result_img = Image.new(img1_pil.mode, (new_width, new_height), background_color)
|
||
result_img.paste(img2_pil, ((new_width - w2) // 2, 0))
|
||
result_img.paste(img1_pil, ((new_width - w1) // 2, h2))
|
||
else: # 理论上不会到这里，因为 direction 是 COMBO
|
||
print(f"错误: {node_name_log} 未知的拼接方向: {direction}")
|
||
continue
|
||
|
||
output_filename = f"frame_{concatenated_count:06d}.png" # 固定输出为png
|
||
output_filepath = os.path.join(final_output_dir, output_filename)
|
||
result_img.save(output_filepath)
|
||
concatenated_count += 1
|
||
|
||
except Exception as e_frame:
|
||
print(f"错误: {node_name_log} 处理帧 {i} (文件: {files1[i]}, {files2[i]}) 时出错: {e_frame}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
# 可以选择跳过此帧或中止
|
||
|
||
finally:
|
||
# 关闭打开的图像文件，以释放资源
|
||
if 'img1_pil' in locals() and hasattr(img1_pil, 'close'): img1_pil.close()
|
||
if 'img2_pil' in locals() and hasattr(img2_pil, 'close'): img2_pil.close()
|
||
if 'img1_orig_pil' in locals() and hasattr(img1_orig_pil, 'close'): img1_orig_pil.close()
|
||
if 'img2_orig_pil' in locals() and hasattr(img2_orig_pil, 'close'): img2_orig_pil.close()
|
||
if 'base_img_pil' in locals() and hasattr(base_img_pil, 'close'): base_img_pil.close()
|
||
if 'target_img_pil' in locals() and hasattr(target_img_pil, 'close'): target_img_pil.close()
|
||
if 'target_img_orig_pil' in locals() and hasattr(target_img_orig_pil, 'close'): target_img_orig_pil.close()
|
||
if 'resized_target_img' in locals() and hasattr(resized_target_img, 'close'): resized_target_img.close()
|
||
if 'result_img' in locals() and hasattr(result_img, 'close'): result_img.close()
|
||
|
||
|
||
comfy_pbar.update(1)
|
||
console_pbar.update(1)
|
||
|
||
end_time_process = time.time()
|
||
print(f"{node_name_log} 拼接完成。总共 {concatenated_count} 帧已保存到 {final_output_dir}")
|
||
print(f"{node_name_log} 方法总执行耗时: {end_time_process - start_time_process:.2f} 秒。")
|
||
|
||
if concatenated_count == 0 and num_frames_to_process > 0:
|
||
return (f"错误: 未成功拼接任何帧，请检查日志。", 0)
|
||
|
||
return (final_output_dir, concatenated_count)
|
||
|
||
# --- ComfyUI 节点注册 ---
|
||
NODE_CLASS_MAPPINGS = {
|
||
"AIIA_Utils_Image_Concanate": AIIA_Utils_Image_Concanate,
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"AIIA_Utils_Image_Concanate": "Image Concatenate (AIIA Utils, Disk)",
|
||
}
|
||
|
||
# 可以在这里添加一个 main 用于独立测试 (如果需要)
|
||
# if __name__ == '__main__':
|
||
# # 创建一些假的目录和图像文件进行测试
|
||
# # ...
|
||
# concatenator = AIIA_Utils_Image_Concanate()
|
||
# # result_dir, count = concatenator.concatenate_images_from_disk(...)
|
||
# # print(f"Test finished. Output to {result_dir}, {count} frames.") |