XieJunchen
5a7ea89870
Implement a node for batch loading multiple file types including images, videos, and audio, with support for various input formats and automatic type classification.
319 lines
11 KiB
Python
319 lines
11 KiB
Python
"""
|
||
多路径文件批量加载节点。
|
||
|
||
用途：一个节点解析用户传入的字符串（多种分隔符），按类型（图片/视频/音频）
|
||
分别输出到固定数量的端口。下游若某个端口无输入需要，直接不连线即可 ——
|
||
本节点对"没有对应文件"的端口返回 None，ComfyUI 不会执行未连线的下游，
|
||
从而保证一个工作流适配 0~N 个素材。
|
||
|
||
默认端口数量与 MiniMaxH3ReferenceToVideo 对齐：9 图 / 3 视频 / 3 音频。
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import tempfile
|
||
|
||
# 重量级依赖延迟到函数内 import，方便离线做纯字符串逻辑的单测
|
||
try:
|
||
import folder_paths # ComfyUI 运行环境注入
|
||
except Exception: # 允许脱离 ComfyUI 时被静态解析
|
||
folder_paths = None
|
||
|
||
|
||
# 与 MiniMax H3 对齐的默认端口数量
|
||
MAX_IMAGES = 9
|
||
MAX_VIDEOS = 3
|
||
MAX_AUDIOS = 3
|
||
|
||
# 支持的扩展名（判断类型用；用户显式把文件填到哪个字段就走哪个字段）
|
||
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff"}
|
||
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm", ".flv", ".m4v"}
|
||
AUDIO_EXTS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus"}
|
||
|
||
# 拆分正则：逗号/分号/竖线/花括号/中英文引号/换行/制表符/多空格
|
||
# 反斜杠 \ 也作为分隔符（用户描述里提到），但注意 Windows 路径本身含 \，
|
||
# 因此优先尝试解析为原始路径，只有当整体不是合法文件时才用 \ 拆一次。
|
||
_SPLIT_PATTERN = re.compile(r"[,;|\n\r\t{}\[\]\"\u201c\u201d\u2018\u2019']+")
|
||
# 看起来像 Windows 绝对路径（C:\..., D:/...）
|
||
_WIN_ABS = re.compile(r"^[A-Za-z]:[\\/]")
|
||
|
||
|
||
def _looks_absolute(p):
|
||
"""跨平台的"看起来是绝对路径"判断：兼顾 POSIX 和 Windows。"""
|
||
return os.path.isabs(p) or bool(_WIN_ABS.match(p))
|
||
|
||
|
||
def _split_paths(text):
|
||
"""把一段用户输入的字符串切成候选路径列表。
|
||
|
||
规则：
|
||
1. 先按正则拆一遍
|
||
2. 每段 strip 掉空白和成对包裹的引号
|
||
3. 单个段落如果不是已存在的文件，再尝试按空格/反斜杠二次拆分
|
||
（这样 "C:\\a.png" 这种 Windows 绝对路径整体不会被误拆）
|
||
"""
|
||
if not text:
|
||
return []
|
||
raw_parts = _SPLIT_PATTERN.split(text)
|
||
out = []
|
||
for p in raw_parts:
|
||
p = p.strip().strip("\"'\u201c\u201d\u2018\u2019")
|
||
if not p:
|
||
continue
|
||
if os.path.isfile(p) or _looks_absolute(p):
|
||
# 整段就是合法文件（或至少是绝对路径），直接收
|
||
out.append(p)
|
||
continue
|
||
# 二次拆：空格和反斜杠。到这里说明段里可能塞了多个相对路径
|
||
for q in re.split(r"[\s\\]+", p):
|
||
q = q.strip()
|
||
if q:
|
||
out.append(q)
|
||
return out
|
||
|
||
|
||
def _resolve(path):
|
||
"""把路径解析成一个真实存在的文件路径。找不到返回 None。
|
||
|
||
顺序：绝对路径 → ComfyUI input 目录相对路径。
|
||
"""
|
||
if not path:
|
||
return None
|
||
if os.path.isfile(path):
|
||
return path
|
||
if folder_paths is not None:
|
||
candidate = os.path.join(folder_paths.get_input_directory(), path)
|
||
if os.path.isfile(candidate):
|
||
return candidate
|
||
return None
|
||
|
||
|
||
def _guess_kind(path):
|
||
ext = os.path.splitext(path)[1].lower()
|
||
if ext in IMAGE_EXTS:
|
||
return "image"
|
||
if ext in VIDEO_EXTS:
|
||
return "video"
|
||
if ext in AUDIO_EXTS:
|
||
return "audio"
|
||
return None
|
||
|
||
|
||
# ---------- 加载器 ----------
|
||
|
||
def _load_image(path):
|
||
"""PIL 加载 → (1, H, W, C) float32 in [0,1]。"""
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image, ImageOps
|
||
with Image.open(path) as im:
|
||
im = ImageOps.exif_transpose(im)
|
||
if im.mode == "RGBA":
|
||
im = im.convert("RGBA")
|
||
elif im.mode != "RGB":
|
||
im = im.convert("RGB")
|
||
arr = np.array(im).astype(np.float32) / 255.0
|
||
if arr.ndim == 2:
|
||
arr = np.expand_dims(arr, -1)
|
||
tensor = torch.from_numpy(arr).unsqueeze(0)
|
||
return tensor
|
||
|
||
|
||
def _extract_audio(video_path):
|
||
"""用 ffmpeg 抽视频音轨为 wav 再读入。没有音轨返回静音 dict。"""
|
||
import torch
|
||
sample_rate = 44100
|
||
silent = {
|
||
"waveform": torch.zeros((1, 2, 1), dtype=torch.float32),
|
||
"sample_rate": sample_rate,
|
||
}
|
||
tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
||
tmp_path = tmp.name
|
||
tmp.close()
|
||
try:
|
||
cmd = [
|
||
"ffmpeg", "-y", "-i", video_path,
|
||
"-vn", "-acodec", "pcm_s16le", "-ar", str(sample_rate), "-ac", "2",
|
||
"-f", "wav", tmp_path,
|
||
]
|
||
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
if not os.path.exists(tmp_path) or os.path.getsize(tmp_path) == 0:
|
||
return silent
|
||
import torchaudio
|
||
waveform, sr = torchaudio.load(tmp_path)
|
||
if waveform.dim() == 2:
|
||
waveform = waveform.unsqueeze(0)
|
||
return {"waveform": waveform, "sample_rate": sr}
|
||
except Exception as e:
|
||
print(f"[MultiFileLoader] 抽取音轨失败，返回静音：{e}")
|
||
return silent
|
||
finally:
|
||
try:
|
||
os.remove(tmp_path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _load_video(path):
|
||
"""加载视频为 (T, H, W, C) 帧序列 + 音轨 AUDIO dict。"""
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
cap = cv2.VideoCapture(path)
|
||
if not cap.isOpened():
|
||
raise FileNotFoundError(f"无法打开视频文件: {path}")
|
||
frames = []
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
frames.append(torch.from_numpy(frame.astype(np.float32) / 255.0))
|
||
cap.release()
|
||
if not frames:
|
||
raise ValueError(f"视频读取到 0 帧: {path}")
|
||
return torch.stack(frames, dim=0), _extract_audio(path)
|
||
|
||
|
||
def _load_audio(path):
|
||
"""加载独立音频文件为 ComfyUI AUDIO dict。"""
|
||
import torchaudio
|
||
waveform, sr = torchaudio.load(path)
|
||
if waveform.dim() == 2:
|
||
waveform = waveform.unsqueeze(0)
|
||
return {"waveform": waveform, "sample_rate": sr}
|
||
|
||
|
||
# ---------- 节点主体 ----------
|
||
|
||
class MultiFileLoader:
|
||
"""一次解析多组路径字符串，按类型分发到固定数量的输出端口。
|
||
|
||
- 未使用的端口返回 None，下游不连线即安全（ComfyUI 不会执行未连线的下游）。
|
||
- 视频端口成对：video_k 输出帧序列 IMAGE，video_audio_k 输出该视频的音轨 AUDIO。
|
||
若视频无音轨则输出静音 dict，避免下游 AudioLoader 报错。
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"optional": {
|
||
"image_paths": ("STRING", {
|
||
"multiline": True, "default": "",
|
||
"tooltip": "图片路径，支持 , ; | 换行 空格 引号 花括号 等分隔",
|
||
}),
|
||
"video_paths": ("STRING", {
|
||
"multiline": True, "default": "",
|
||
"tooltip": "视频路径，同上分隔规则",
|
||
}),
|
||
"audio_paths": ("STRING", {
|
||
"multiline": True, "default": "",
|
||
"tooltip": "音频路径，同上分隔规则",
|
||
}),
|
||
"auto_classify": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "开启时会按扩展名自动纠正类型（比如把 mp4 从 image_paths 归到视频输出）",
|
||
}),
|
||
}
|
||
}
|
||
|
||
# 动态生成 RETURN_TYPES/RETURN_NAMES
|
||
RETURN_TYPES = (
|
||
*(("IMAGE",) * MAX_IMAGES),
|
||
*(("IMAGE", "AUDIO") * MAX_VIDEOS),
|
||
*(("AUDIO",) * MAX_AUDIOS),
|
||
"INT", "INT", "INT",
|
||
)
|
||
RETURN_NAMES = (
|
||
*(f"image_{i + 1}" for i in range(MAX_IMAGES)),
|
||
*(name for i in range(MAX_VIDEOS) for name in (f"video_{i + 1}", f"video_audio_{i + 1}")),
|
||
*(f"audio_{i + 1}" for i in range(MAX_AUDIOS)),
|
||
"image_count", "video_count", "audio_count",
|
||
)
|
||
FUNCTION = "load"
|
||
CATEGORY = "云服务/多路加载"
|
||
OUTPUT_NODE = False
|
||
|
||
def load(self, image_paths="", video_paths="", audio_paths="", auto_classify=True):
|
||
# 1. 拆分三组字符串
|
||
raw = {
|
||
"image": _split_paths(image_paths),
|
||
"video": _split_paths(video_paths),
|
||
"audio": _split_paths(audio_paths),
|
||
}
|
||
|
||
# 2. 解析真实路径 + 类型分桶
|
||
bucket = {"image": [], "video": [], "audio": []}
|
||
for declared_kind, paths in raw.items():
|
||
for p in paths:
|
||
resolved = _resolve(p)
|
||
if resolved is None:
|
||
print(f"[MultiFileLoader] 找不到文件，跳过：{p}")
|
||
continue
|
||
kind = declared_kind
|
||
if auto_classify:
|
||
guessed = _guess_kind(resolved)
|
||
if guessed is not None and guessed != declared_kind:
|
||
print(f"[MultiFileLoader] 自动纠正类型 {declared_kind}→{guessed}: {resolved}")
|
||
kind = guessed
|
||
bucket[kind].append(resolved)
|
||
|
||
# 3. 数量校验（超出直接截断，给个 warning）
|
||
if len(bucket["image"]) > MAX_IMAGES:
|
||
print(f"[MultiFileLoader] 图片超过上限 {MAX_IMAGES}，多余的已忽略")
|
||
bucket["image"] = bucket["image"][:MAX_IMAGES]
|
||
if len(bucket["video"]) > MAX_VIDEOS:
|
||
print(f"[MultiFileLoader] 视频超过上限 {MAX_VIDEOS}，多余的已忽略")
|
||
bucket["video"] = bucket["video"][:MAX_VIDEOS]
|
||
if len(bucket["audio"]) > MAX_AUDIOS:
|
||
print(f"[MultiFileLoader] 音频超过上限 {MAX_AUDIOS}，多余的已忽略")
|
||
bucket["audio"] = bucket["audio"][:MAX_AUDIOS]
|
||
|
||
# 4. 逐个加载，未填满的端口填 None
|
||
image_outs = [None] * MAX_IMAGES
|
||
for i, p in enumerate(bucket["image"]):
|
||
try:
|
||
image_outs[i] = _load_image(p)
|
||
except Exception as e:
|
||
print(f"[MultiFileLoader] 加载图片失败 {p}: {e}")
|
||
|
||
video_frame_outs = [None] * MAX_VIDEOS
|
||
video_audio_outs = [None] * MAX_VIDEOS
|
||
for i, p in enumerate(bucket["video"]):
|
||
try:
|
||
frames, audio = _load_video(p)
|
||
video_frame_outs[i] = frames
|
||
video_audio_outs[i] = audio
|
||
except Exception as e:
|
||
print(f"[MultiFileLoader] 加载视频失败 {p}: {e}")
|
||
|
||
audio_outs = [None] * MAX_AUDIOS
|
||
for i, p in enumerate(bucket["audio"]):
|
||
try:
|
||
audio_outs[i] = _load_audio(p)
|
||
except Exception as e:
|
||
print(f"[MultiFileLoader] 加载音频失败 {p}: {e}")
|
||
|
||
# 5. 按 RETURN_NAMES 顺序拼输出
|
||
result: list = list(image_outs)
|
||
for f, a in zip(video_frame_outs, video_audio_outs):
|
||
result.append(f)
|
||
result.append(a)
|
||
result.extend(audio_outs)
|
||
result.extend([
|
||
sum(1 for x in image_outs if x is not None),
|
||
sum(1 for x in video_frame_outs if x is not None),
|
||
sum(1 for x in audio_outs if x is not None),
|
||
])
|
||
return tuple(result)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MultiFileLoader": MultiFileLoader,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MultiFileLoader": "📦 多路径文件批量加载",
|
||
}
|