diff --git a/__init__.py b/__init__.py index 28522b1..bd99d6f 100644 --- a/__init__.py +++ b/__init__.py @@ -6,5 +6,5 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "BatchPromptReaderWithClip": "批量提示词加载器 (Batch Prompt Loader)", + "BatchPromptReaderWithClip": "Batch Prompt Loader", } diff --git a/batch_prompt_loader.py b/batch_prompt_loader.py index ffb28db..83d7037 100644 --- a/batch_prompt_loader.py +++ b/batch_prompt_loader.py @@ -3,6 +3,9 @@ ComfyUI Batch Prompt Loader - Enhanced Version A custom node for loading and encoding prompts from TXT files with advanced features: - Recursive scanning, wildcard filtering, skip exists, metadata output, preview display - External log recording: records which prompt files have been processed, and auto-skips them next run + +All user-supplied paths (folder_path / log_folder / output_folder) are confined to the +ComfyUI root directory to prevent path traversal. """ import os import torch @@ -15,8 +18,28 @@ _global_mode_state = {} _log_cleared_this_session = False +def _confine_to_base(base_dir, user_path, label): + """Resolve a user-supplied path and verify it stays inside base_dir. + + Mitigates path traversal: absolute paths pointing outside the ComfyUI root and + relative paths escaping via '..' are rejected. Symlinks are resolved via realpath. + """ + base_real = os.path.normcase(os.path.realpath(base_dir)) + if os.path.isabs(user_path): + candidate = user_path + else: + candidate = os.path.join(base_dir, user_path) + candidate_real = os.path.normcase(os.path.realpath(candidate)) + if candidate_real != base_real and not candidate_real.startswith(base_real + os.sep): + raise Exception( + f"[BatchPromptLoader] '{label}' must resolve to a location inside the ComfyUI directory.\n" + f"Resolved path: {candidate_real}\nAllowed base: {base_real}" + ) + return candidate_real + + def _sanitize_log_name(rel_path): - """将相对路径转为安全的日志文件名（替换路径分隔符和非法字符）""" + """Convert a relative prompt path into a safe log filename (replacing separators and illegal chars).""" name = rel_path.replace('\\', '_').replace('/', '_') for ch in '<>:"/\\|?*': name = name.replace(ch, '_') @@ -24,7 +47,7 @@ def _sanitize_log_name(rel_path): def _load_log(log_dir, source_dir): - """读取日志文件夹中匹配当前来源目录的已处理文件集合""" + """Read the set of already-processed prompt files matching the current source directory.""" processed = set() if not os.path.isdir(log_dir): return processed @@ -40,12 +63,12 @@ def _load_log(log_dir, source_dir): except Exception: continue except Exception as e: - print(f"[BatchPromptLoader] Warning: 读取日志失败: {e}") + print(f"[BatchPromptLoader] Warning: failed to read log: {e}") return processed def _write_log(log_dir, source_dir, prompt_file, index, total_count): - """将已处理的提示词文件写入日志（每个文件一个 JSON，便于选择性删除）""" + """Record a processed prompt file (one JSON per file, so entries can be deleted selectively).""" try: os.makedirs(log_dir, exist_ok=True) log_name = _sanitize_log_name(prompt_file) + '.json' @@ -59,7 +82,8 @@ def _write_log(log_dir, source_dir, prompt_file, index, total_count): with open(os.path.join(log_dir, log_name), 'w', encoding='utf-8') as f: json.dump(entry, f, ensure_ascii=False, indent=2) except Exception as e: - print(f"[BatchPromptLoader] Warning: 写入日志失败: {e}") + print(f"[BatchPromptLoader] Warning: failed to write log: {e}") + class BatchPromptReaderWithClip: @classmethod @@ -70,7 +94,7 @@ class BatchPromptReaderWithClip: "folder_path": ("STRING", { "default": "input/batch_prompts", "multiline": False, - "label": "文件夹路径" + "label": "Folder Path" }), "current_number": ("INT", { "default": 0, @@ -78,44 +102,44 @@ class BatchPromptReaderWithClip: "max": 999999, "step": 1, "control_after_generate": ["fixed", "increment", "decrement", "random"], - "label": "当前编号" + "label": "Current Index" }), "recursive": ("BOOLEAN", { "default": True, - "label": "递归扫描子文件夹" + "label": "Recursive Scan" }), "reverse_order": ("BOOLEAN", { "default": False, - "label": "倒序排列文件" + "label": "Reverse Order" }), "file_pattern": ("STRING", { "default": "*.txt", "multiline": False, - "label": "文件名过滤 (通配符)" + "label": "Filename Pattern (wildcard)" }), "output_folder": ("STRING", { "default": "output/", "multiline": False, - "label": "输出目录 (用于跳过检查)" + "label": "Output Folder (skip check)" }), "enable_logging": ("BOOLEAN", { "default": False, - "label": "启用日志记录 (跳过已处理)" + "label": "Enable Logging (skip processed)" }), "log_folder": ("STRING", { "default": "user/default/batch_prompt_logs", "multiline": False, - "label": "日志文件夹 (外置可清除)" + "label": "Log Folder (external, clearable)" }), "clear_log_on_start": ("BOOLEAN", { "default": False, - "label": "会话启动时清空日志" + "label": "Clear Log On Session Start" }), }, "optional": { "skip_exists": ("BOOLEAN", { "default": False, - "label": "跳过已存在的图片" + "label": "Skip Existing Images" }), }, "hidden": { @@ -123,112 +147,118 @@ class BatchPromptReaderWithClip: "unique_id": "UNIQUE_ID" } } - + RETURN_TYPES = ("CONDITIONING", "STRING", "INT", "INT") RETURN_NAMES = ("conditioning", "filename", "index", "total_count") FUNCTION = "read_and_encode" OUTPUT_NODE = True CATEGORY = "batch_tools" - - def read_and_encode(self, clip, folder_path, current_number, recursive=True, reverse_order=False, + + def read_and_encode(self, clip, folder_path, current_number, recursive=True, reverse_order=False, file_pattern="*.txt", output_folder="output/", skip_exists=False, enable_logging=False, log_folder="user/default/batch_prompt_logs", clear_log_on_start=False, extra_pnginfo=None, unique_id=None): base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - - if os.path.isabs(folder_path): - target_dir = folder_path - else: - target_dir = os.path.join(base_dir, folder_path) - - # 优化 3: 增强错误提示 - 先检查路径是否存在 + + # Confine user-supplied paths to the ComfyUI root (path traversal protection) + target_dir = _confine_to_base(base_dir, folder_path, "folder_path") + if not os.path.exists(target_dir): - raise Exception(f"[错误] 路径不存在：{target_dir}\n\n请检查 folder_path 是否正确。") - + raise Exception( + f"[BatchPromptLoader] Folder does not exist: {target_dir}\n\n" + f"Please check the 'folder_path' value." + ) + if not os.path.isdir(target_dir): - raise Exception(f"[错误] 路径不是文件夹：{target_dir}\n\n请确保这是一个有效的目录。") - + raise Exception( + f"[BatchPromptLoader] Not a directory: {target_dir}\n\n" + f"Please make sure this is a valid folder." + ) + try: txt_files = [] - - # 解析文件模式 + pattern_lower = file_pattern.lower().strip() if not pattern_lower or pattern_lower == '*.txt': search_pattern = '*.txt' required_ext = '.txt' else: search_pattern = pattern_lower - # 提取扩展名（如 *.prompt.txt -> .prompt.txt） if '*' in search_pattern: ext_part = search_pattern.split('*', 1)[-1] required_ext = ext_part if ext_part else '.txt' else: required_ext = '.txt' - + if recursive: - # 递归扫描所有子文件夹 for root, dirs, files in os.walk(target_dir): for f in files: - # 优化 2: 通配符过滤 if fnmatch.fnmatch(f.lower(), search_pattern.lower()): if required_ext and not f.lower().endswith(required_ext): continue rel_path = os.path.relpath(os.path.join(root, f), target_dir) txt_files.append(rel_path) else: - # 只读取根目录 all_files = os.listdir(target_dir) for f in all_files: if fnmatch.fnmatch(f.lower(), search_pattern.lower()): if required_ext and not f.lower().endswith(required_ext): continue txt_files.append(f) - - # 排序 + txt_files.sort(key=lambda x: x.lower()) - - # 倒序处理 + if reverse_order: txt_files.reverse() - + except Exception as e: - raise Exception(f"[错误] 读取文件夹失败：{target_dir}\n\n错误信息：{str(e)}") - - # 优化 3: 空文件夹检测 + raise Exception(f"[BatchPromptLoader] Failed to read folder: {target_dir}\n\nError: {str(e)}") + if len(txt_files) == 0: - raise Exception(f"[错误] 文件夹中没有匹配的文件：{target_dir}\n\n搜索模式：{search_pattern}\n\n请添加对应的 TXT 文件或调整文件名过滤规则。") - - # 优化 6: 外置日志记录 - 跳过已处理过的提示词文件 + raise Exception( + f"[BatchPromptLoader] No matching files in folder: {target_dir}\n\n" + f"Search pattern: {search_pattern}\n\n" + f"Please add matching TXT files or adjust the filename pattern." + ) + + # External logging: skip already-processed prompt files log_dir = None if enable_logging: - if os.path.isabs(log_folder): - log_dir = log_folder - else: - log_dir = os.path.join(base_dir, log_folder) - - # 会话启动后第一次运行时清空日志（配合 clear_log_on_start） + log_dir = _confine_to_base(base_dir, log_folder, "log_folder") + global _log_cleared_this_session if clear_log_on_start and not _log_cleared_this_session: try: if os.path.isdir(log_dir): for fn in os.listdir(log_dir): - if fn.endswith('.json'): - os.remove(os.path.join(log_dir, fn)) + if not fn.endswith('.json'): + continue + # Only remove log files previously written by this node + log_path = os.path.join(log_dir, fn) + try: + with open(log_path, 'r', encoding='utf-8') as lf: + data = json.load(lf) + if isinstance(data, dict) and 'prompt_file' in data: + os.remove(log_path) + except Exception: + continue _log_cleared_this_session = True - print("[BatchPromptLoader] 🧹 已清空日志文件夹（会话启动后第一次运行）") + print("[BatchPromptLoader] Cleared log folder (first run this session)") except Exception as e: - print(f"[BatchPromptLoader] Warning: 清空日志失败: {e}") - + print(f"[BatchPromptLoader] Warning: failed to clear log: {e}") + processed_files = _load_log(log_dir, os.path.normpath(target_dir)) if processed_files: remaining = [f for f in txt_files if f not in processed_files] - print(f"[BatchPromptLoader] 📋 日志过滤：已跳过 {len(txt_files) - len(remaining)} 个已处理文件，剩余 {len(remaining)} 个待处理") + print(f"[BatchPromptLoader] Log filter: skipped {len(txt_files) - len(remaining)} " + f"processed file(s), {len(remaining)} remaining") if not remaining: - print("[BatchPromptLoader] ✅ 所有提示词已处理完成（清除日志文件夹可重新开始）") + print("[BatchPromptLoader] All prompts already processed " + "(clear the log folder to start over)") empty_cond = [[torch.zeros(1, 1, 2816), {"pooled_output": torch.zeros(1, 2816)}]] return (empty_cond, "ALL_DONE", 0, 0) txt_files = remaining - + counter_file = os.path.join(base_dir, "user", "default", "batch_prompt_counter.json") try: if os.path.exists(counter_file) and os.path.getsize(counter_file) > 0: @@ -236,91 +266,94 @@ class BatchPromptReaderWithClip: counter_data = json.load(f) else: counter_data = {"last_index": 0} - except: + except Exception: counter_data = {"last_index": 0} - + mode = "fixed" effective_current_number = current_number - + if isinstance(current_number, (list, tuple)) and len(current_number) >= 2: effective_current_number = current_number[0] mode = current_number[1] elif isinstance(current_number, dict) and 'mode' in current_number: mode = current_number['mode'] effective_current_number = current_number.get('value', 0) - + if mode == "increment": old_index = counter_data.get("last_index", effective_current_number) new_index = (old_index + 1) % len(txt_files) effective_index = new_index counter_data["last_index"] = new_index - print(f"[BatchPromptLoader] Increment: {old_index} → {new_index} ({txt_files[new_index][:30]}...)") - + print(f"[BatchPromptLoader] Increment: {old_index} -> {new_index} " + f"({txt_files[new_index][:30]}...)") + elif mode == "decrement": old_index = counter_data.get("last_index", effective_current_number) new_index = (old_index - 1) % len(txt_files) effective_index = new_index counter_data["last_index"] = new_index - print(f"[BatchPromptLoader] Decrement: {old_index} → {new_index} ({txt_files[new_index][:30]}...)") - + print(f"[BatchPromptLoader] Decrement: {old_index} -> {new_index} " + f"({txt_files[new_index][:30]}...)") + elif mode == "random": effective_index = random_module.randint(0, len(txt_files) - 1) counter_data["last_index"] = effective_index - print(f"[BatchPromptLoader] Random: {effective_index} ({txt_files[effective_index][:30]}...)") - + print(f"[BatchPromptLoader] Random: {effective_index} " + f"({txt_files[effective_index][:30]}...)") + else: effective_index = effective_current_number % len(txt_files) counter_data["last_index"] = effective_index - print(f"[BatchPromptLoader] Fixed: {effective_index} ({txt_files[effective_index][:30]}...)") - + print(f"[BatchPromptLoader] Fixed: {effective_index} " + f"({txt_files[effective_index][:30]}...)") + try: os.makedirs(os.path.dirname(counter_file), exist_ok=True) with open(counter_file, 'w', encoding='utf-8') as f: json.dump(counter_data, f, ensure_ascii=False) except Exception as e: print(f"[BatchPromptLoader] Warning: Failed to save counter: {e}") - + selected_file = txt_files[effective_index] file_path = os.path.join(target_dir, selected_file) - - # 优化 4: 跳过已存在的图片 + + # Skip images that already exist in the output folder if skip_exists and output_folder.strip(): - out_base = os.path.join(base_dir, output_folder.rstrip('/').rstrip('\\')) + out_base = _confine_to_base(base_dir, output_folder.rstrip('/').rstrip('\\'), "output_folder") base_name = os.path.splitext(os.path.basename(selected_file))[0] possible_extensions = ['.png', '.jpg', '.jpeg', '.webp'] file_exists = False - + for ext in possible_extensions: check_path = os.path.join(out_base, f"{base_name}{ext}") if os.path.exists(check_path): file_exists = True - print(f"[BatchPromptLoader] ⏭️ 跳过已存在：{base_name}{ext}") + print(f"[BatchPromptLoader] Skipped existing: {base_name}{ext}") break - + if file_exists: - print(f"[BatchPromptLoader] 已跳过：{selected_file}") if enable_logging: - _write_log(log_dir, os.path.normpath(target_dir), selected_file, effective_index, len(txt_files)) + _write_log(log_dir, os.path.normpath(target_dir), selected_file, + effective_index, len(txt_files)) empty_cond = [[torch.zeros(1, 1, 2816), {"pooled_output": torch.zeros(1, 2816)}]] return (empty_cond, f"SKIPPED:{selected_file}", effective_index, len(txt_files)) - - # 优化 1: 显示预览信息 - preview_msg = f"✓ [{effective_index + 1}/{len(txt_files)}] {selected_file}" + + preview_msg = f"[{effective_index + 1}/{len(txt_files)}] {selected_file}" print(f"[BatchPromptLoader] {preview_msg}") - + try: with open(file_path, 'r', encoding='utf-8') as f: content = f.read() - + positive = "" negative = "" section = None - + for line in content.split('\n'): line = line.strip() if not line: continue - + lower_line = line.lower() if lower_line.startswith('positive:'): section = 'positive' @@ -332,47 +365,46 @@ class BatchPromptReaderWithClip: positive += (' ' if positive else '') + line elif section == 'negative': negative += (' ' if negative else '') + line - + if not positive and content.strip(): positive = content.strip() if not negative: negative = "" - + with torch.no_grad(): tokens_positive = clip.tokenize(positive) tokens_negative = clip.tokenize(negative) - + cond_positive, pooled_positive = clip.encode_from_tokens(tokens_positive, return_pooled=True) cond_negative, pooled_negative = clip.encode_from_tokens(tokens_negative, return_pooled=True) - + if cond_positive.dim() == 2: cond_positive = cond_positive.unsqueeze(0) if cond_negative.dim() == 2: cond_negative = cond_negative.unsqueeze(0) - + if pooled_positive is None: pooled_positive = torch.zeros(1, 2816) elif pooled_positive.dim() == 1: pooled_positive = pooled_positive.unsqueeze(0) - + if pooled_negative is None: pooled_negative = torch.zeros(1, 2816) elif pooled_negative.dim() == 1: pooled_negative = pooled_negative.unsqueeze(0) - + conditioning = [[cond_positive, {"pooled_output": pooled_positive}]] - - # 优化 6: 记录到外置日志（成功编码后标记为已处理） + if enable_logging: - _write_log(log_dir, os.path.normpath(target_dir), selected_file, effective_index, len(txt_files)) - - # 优化 5: 输出元数据 + _write_log(log_dir, os.path.normpath(target_dir), selected_file, + effective_index, len(txt_files)) + return ( conditioning, selected_file, # filename - effective_index, # index (从 0 开始) + effective_index, # index (0-based) len(txt_files) # total_count ) - + except Exception as e: raise Exception(f"Failed to read file: {selected_file}\nError: {str(e)}")