From f17aa7db8a5e7f7bf913352291758b1c3b7ab545 Mon Sep 17 00:00:00 2001 From: kjqwer <2990346238@qq.com> Date: Thu, 23 Jul 2026 13:07:49 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=9E=B6=E6=9E=84=EF=BC=8C?= =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=8F=AF=E8=A7=86=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- MANIFEST.in | 4 + README.md | 228 +------- __init__.py | 21 +- locales/en/main.json | 3 +- locales/en/nodeDefs.json | 218 +------- locales/zh/main.json | 3 +- locales/zh/nodeDefs.json | 218 +------- pyproject.toml | 4 +- requirements.txt | 1 - src/SmartSaveImage/__init__.py | 8 + src/SmartSaveImage/core/__init__.py | 12 - src/SmartSaveImage/core/image_utils.py | 168 ------ src/SmartSaveImage/core/metadata.py | 201 ------- src/SmartSaveImage/core/path_utils.py | 230 -------- src/SmartSaveImage/nodes.py | 118 ---- src/SmartSaveImage/nodes/__init__.py | 14 +- src/SmartSaveImage/nodes/folder_manager.py | 412 -------------- src/SmartSaveImage/nodes/image_saver.py | 332 ----------- src/SmartSaveImage/nodes/smart_save.py | 605 +++++++++++++++++++++ src/SmartSaveImage/server_routes.py | 125 +++++ src/SmartSaveImage/utils/__init__.py | 7 - src/SmartSaveImage/utils/validators.py | 147 ----- tests/pytest.ini | 6 +- tests/test_SmartSaveImage.py | 200 ++++++- web/smart_save.css | 270 +++++++++ web/smart_save.js | 445 +++++++++++++++ 27 files changed, 1725 insertions(+), 2278 deletions(-) delete mode 100644 src/SmartSaveImage/core/__init__.py delete mode 100644 src/SmartSaveImage/core/image_utils.py delete mode 100644 src/SmartSaveImage/core/metadata.py delete mode 100644 src/SmartSaveImage/core/path_utils.py delete mode 100644 src/SmartSaveImage/nodes.py delete mode 100644 src/SmartSaveImage/nodes/folder_manager.py delete mode 100644 src/SmartSaveImage/nodes/image_saver.py create mode 100644 src/SmartSaveImage/nodes/smart_save.py create mode 100644 src/SmartSaveImage/server_routes.py delete mode 100644 src/SmartSaveImage/utils/__init__.py delete mode 100644 src/SmartSaveImage/utils/validators.py create mode 100644 web/smart_save.css create mode 100644 web/smart_save.js diff --git a/.gitignore b/.gitignore index 246b1d7..f047fb2 100644 --- a/.gitignore +++ b/.gitignore @@ -125,4 +125,5 @@ temp/ output/ models/ input/ -test/ \ No newline at end of file +test/ +tests/ diff --git a/MANIFEST.in b/MANIFEST.in index 6514960..98014c3 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,10 @@ include LICENSE include README.md +recursive-include src *.py +recursive-include web *.js *.css +recursive-include locales *.json + recursive-exclude * __pycache__ recursive-exclude * *.py[co] diff --git a/README.md b/README.md index 5b9e9dc..37beea5 100644 --- a/README.md +++ b/README.md @@ -1,209 +1,37 @@ -# SmartSaveImage - 智能图片保存节点 +# SmartSaveImage -[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) -[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/) -[![ComfyUI](https://img.shields.io/badge/ComfyUI-Compatible-green.svg)](https://github.com/comfyanonymous/ComfyUI) +一个带实时路径预览的 ComfyUI 图片保存节点。节点菜单中搜索 **智能保存图片**。 -一个功能强大的ComfyUI自定义节点包，提供智能的文件夹管理和图片保存功能。 +## 使用方式 -## 🚀 安装 +1. 连接 `images`。 +2. 选择保存位置；自定义模式可填写绝对路径或相对 `output` 的路径。 +3. 填写子目录规则和文件名规则。 +4. 在“保存结果预览”中确认最终目录与示例文件名。 -### 方法一：通过ComfyUI Manager安装（推荐） -1. 打开ComfyUI Manager -2. 搜索 "SmartSaveImage" -3. 点击安装 +模板变量默认折叠，展开后点击变量即可插入当前规则输入框。 -### 方法二：手动安装 -```bash -cd ComfyUI/custom_nodes -git clone https://github.com/kjqwer/SmartSaveImage.git -cd SmartSaveImage -pip install -r requirements.txt +## 模板变量 + +- 时间：`%date:yyyy-MM-dd%`、`%year%`、`%month%`、`%day%`、`%hour%`、`%minute%`、`%second%` +- 模型：`%model%`、`%model_full%`、`%unet%`、`%lora%`、`%vae%` +- 采样：`%seed%`、`%steps%`、`%cfg%`、`%sampler%`、`%scheduler%` +- 图片：`%width%`、`%height%`、`%prompt%`、`%batch%` + +例如，保存到自定义根目录下的模型文件夹： + +```text +子目录规则：krea/%model% +文件名规则：image +同名冲突：自动编号 ``` -### 方法三：下载ZIP -1. 下载本仓库的ZIP文件 -2. 解压到 `ComfyUI/custom_nodes/` 目录 -3. 重启ComfyUI +## 保存行为 -## 🌟 主要特性 +- 支持 PNG、JPEG、WebP；PNG 默认压缩等级为 4，与 ComfyUI 自带保存节点一致，可在 0-9 间调整且始终无损；WebP 使用无损保存。 +- 自动编号不会覆盖已有文件；覆盖模式也会保证批量图片互不覆盖。 +- PNG 使用 ComfyUI 原生元数据字段；关闭“嵌入工作流”后不写入生成信息。 +- 保存到 `output`、`input`、`temp` 时直接使用对应预览类型；其他绝对目录使用临时预览。 +- 模型名自动从工作流读取，也可在“模型来源”中手动指定。 -- **智能文件夹管理** - 自动创建有组织的文件夹结构 -- **灵活的保存选项** - 支持多种图片格式和质量设置 -- **元数据嵌入** - 自动提取并保存工作流信息 -- **批量处理** - 高效处理多张图片 -- **用户友好** - 直观的界面和丰富的选项 - -## 📦 节点介绍 - -### 智能文件夹管理器 (SmartFolderManager) -负责创建和管理文件夹结构，从工作流中自动提取元数据。 - -### 智能图片保存器 (SmartImageSaver) -负责保存图片，支持多种格式、压缩选项和元数据嵌入。 - -## 🚀 快速开始 - -### 基本使用流程 - -1. **添加智能文件夹管理器节点** - - 在ComfyUI中搜索"智能文件夹管理器" - - 将要保存的图片连接到管理器的images输入 - -2. **配置文件夹结构** - - 使用开关控制各层文件夹：日期、模型、种子、提示词、自定义 - - 设置基础文件夹路径 - - 可选择连接外部节点（模型、条件、潜在空间）获取更多信息 - -3. **添加智能图片保存器节点** - - 搜索"智能图片保存器" - - 将文件夹管理器的三个输出全部连接到保存器： - - images → images - - folder_path → folder_path - - metadata_json → metadata_json - -4. **配置保存选项** - - 选择文件格式和质量设置 - - 设置文件名和预览模式 - -## 📁 文件夹结构配置 - -### 灵活的层级控制 -现在可以通过开关独立控制每一层文件夹的创建： - -- **日期文件夹** (`enable_date_folder`) - - 按日期组织：`2024-11-15/` - - 可自定义格式：`yyyy-MM-dd`, `yyyy/MM/dd` 等 - - 可选择包含时间：`2024-11-15_14-30-25/` - -- **模型文件夹** (`enable_model_folder`) - - 按模型组织：`sdxl_base/` - - 自动从工作流提取模型名称 - - 支持手动指定或从模型节点输入 - -- **种子文件夹** (`enable_seed_folder`) - - 按种子组织：`seed_12345/` - - 自动从工作流提取种子值 - - 支持手动设置种子 - -- **提示词文件夹** (`enable_prompt_folder`) - - 按提示词组织：`beautiful_landscape/` - - 可设置最大长度，自动清理非法字符 - - 支持手动输入或从条件节点获取 - -- **自定义文件夹** (`enable_custom_folder`) - - 完全自定义：`my_project/` - - 可以是任意文件夹名称 - -### 组合示例 -- 全开：`2024-11-15/sdxl_base/seed_12345/beautiful_landscape/my_project/` -- 仅日期+模型：`2024-11-15/sdxl_base/` -- 仅种子+自定义：`seed_12345/experiment_01/` - -## 🖼️ 图片保存选项 - -### 文件格式支持 -- **PNG** - 无损压缩，支持透明度 -- **JPEG** - 有损压缩，文件较小 -- **WebP** - 现代格式，支持无损和有损 -- **BMP** - 位图格式 -- **TIFF** - 高质量格式 - -### 质量设置 -- **JPEG质量**：1-100（推荐95） -- **WebP质量**：1-100（推荐90） -- **WebP无损**：启用无损压缩 -- **PNG压缩**：0-9级别（推荐6） - -### 文件命名选项 -- **文件名前缀**：自定义前缀 -- **添加时间戳**：在文件名中包含时间 -- **添加计数器**：多张图片时的序号 -- **计数器设置**：起始值和位数 - -## 🔧 高级功能 - -### 智能元数据获取 -- **外部节点优先**：连接外部节点时优先从节点获取信息 -- **工作流自动提取**：没有外部输入时从工作流中自动提取 -- **图片尺寸检测**：直接从图片数据中获取准确尺寸 -- **手动补充**：仅在需要时手动输入补充信息 - -### 元数据嵌入 -- **参数记录**：保存采样器、CFG、步数等技术参数 -- **工作流保存**：可选择嵌入完整工作流信息 -- **多格式支持**：PNG使用PngInfo，JPEG/WebP使用EXIF - -### 预览模式 -- **保存并预览**：保存文件同时在界面显示 -- **仅预览**：只在界面显示，不保存文件 -- **仅保存**：只保存文件，不显示预览 - -### 文件管理 -- **覆盖保护**：避免意外覆盖现有文件 -- **自动重命名**：文件冲突时自动生成新名称 -- **备份功能**：覆盖前创建备份文件 - -## 💡 使用技巧 - -### 推荐工作流设置 - -1. **日常使用** - - 开启：日期文件夹 + 模型文件夹 - - 文件格式：PNG（质量优先）或WebP（体积优先） - - 预览模式：保存并预览 - -2. **批量实验** - - 开启：日期文件夹 + 种子文件夹 + 自定义文件夹 - - 启用计数器，使用描述性前缀 - - 考虑JPEG格式节省空间 - -3. **项目管理** - - 开启：自定义文件夹（项目名）+ 模型文件夹 + 种子文件夹 - - 手动设置模型和种子确保一致性 - - 嵌入完整元数据便于追溯 - -4. **连接外部节点获取准确信息** - - 将CheckpointLoader的MODEL输出连接到model_input - - 将CLIPTextEncode的CONDITIONING连接到conditioning_positive/negative - - 将KSampler的LATENT输出连接到latent_input - - 连接后会优先使用外部节点的信息，无需手动选择模型 - -### 文件夹路径设置 - -- **相对路径**：基于ComfyUI输出目录 - - `output` 或留空 → 默认输出目录 - - `my_project` → `ComfyUI/output/my_project/` - -- **绝对路径**：指定完整路径 - - `D:/AI_Images/` → 直接保存到指定位置 - -### 元数据利用 - -生成的图片会包含丰富的元数据信息： -- 在图片查看器中可以看到生成参数 -- 便于后续复现相同效果 -- 支持批量分析和管理 - -## ⚠️ 注意事项 - -1. **路径权限**：确保ComfyUI对目标文件夹有写入权限 -2. **文件名长度**：避免过长的文件名（建议<200字符） -3. **特殊字符**：文件名会自动清理非法字符 -4. **磁盘空间**：注意监控存储空间，特别是使用无损格式时 - -## 🔄 更新日志 - -### v0.1.0 -- 重新设计的模块化架构 -- 改进的元数据提取功能 -- 更灵活的文件夹管理选项 -- 增强的错误处理和用户反馈 - -## 🤝 支持与反馈 - -如果遇到问题或有改进建议，欢迎反馈！ - ---- - -*享受更智能的图片保存体验！* 🎨 \ No newline at end of file +插件只依赖 ComfyUI 已包含的 Pillow 与 NumPy，不需要额外配置环境。 diff --git a/__init__.py b/__init__.py index 9688bde..10657cd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,19 +1,16 @@ -"""智能保存图片 - ComfyUI节点包""" +"""智能保存图片 - ComfyUI 节点包""" -import os +__author__ = "kj" +__email__ = "2990346238@qq.com" +__version__ = "2.0.0" + +from .src.SmartSaveImage import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# 前端资源目录（可视化面板 JS/CSS） +WEB_DIRECTORY = "./web" __all__ = [ "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY", ] - -__author__ = "kj" -__email__ = "2990346238@qq.com" -__version__ = "0.1.0" - -# 从子包导入节点映射 -from .src.SmartSaveImage import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - -# 声明Web目录以支持国际化 - 指向当前目录让ComfyUI能找到locales文件夹 -WEB_DIRECTORY = os.path.dirname(__file__) \ No newline at end of file diff --git a/locales/en/main.json b/locales/en/main.json index 987a0a0..8cecdba 100644 --- a/locales/en/main.json +++ b/locales/en/main.json @@ -1,4 +1,3 @@ { - "SmartFolderManager": "Smart Folder Manager", - "SmartImageSaver": "Smart Image Saver" + "SmartSaveImage": "Smart Save Image" } diff --git a/locales/en/nodeDefs.json b/locales/en/nodeDefs.json index 0e04d95..a183227 100644 --- a/locales/en/nodeDefs.json +++ b/locales/en/nodeDefs.json @@ -1,206 +1,20 @@ { - "SmartFolderManager": { - "display_name": "Smart Folder Manager", - "description": "Smart folder manager - manages folder structure and metadata collection", + "SmartSaveImage": { + "display_name": "Smart Save Image", + "description": "Preview the destination and organize folders and filenames with templates", "inputs": { - "images": { - "name": "Images", - "tooltip": "Input images for extracting size info and passing to save node" - }, - "base_folder": { - "name": "Base Folder", - "tooltip": "Base folder path, can be relative or absolute path" - }, - "create_subfolders": { - "name": "Create Subfolders", - "tooltip": "Whether to automatically create subfolders" - }, - "enable_date_folder": { - "name": "Enable Date Folder", - "tooltip": "Whether to create date folder" - }, - "enable_model_folder": { - "name": "Enable Model Folder", - "tooltip": "Whether to create model folder" - }, - "enable_seed_folder": { - "name": "Enable Seed Folder", - "tooltip": "Whether to create seed folder" - }, - "enable_prompt_folder": { - "name": "Enable Prompt Folder", - "tooltip": "Whether to create prompt folder" - }, - "enable_custom_folder": { - "name": "Enable Custom Folder", - "tooltip": "Whether to use custom folder" - }, - "date_format": { - "name": "Date Format", - "tooltip": "Date format: yyyy=year MM=month dd=day hh=hour mm=minute ss=second" - }, - "include_time": { - "name": "Include Time", - "tooltip": "Whether to include time in date" - }, - "model_source": { - "name": "Model Source", - "tooltip": "Model info source: auto=from metadata/external input, manual=manual selection", - "options": { - "auto": "Auto Detect", - "manual": "Manual Select" - } - }, - "manual_model_name": { - "name": "Manual Model Name", - "tooltip": "Manual model selection (only effective when model_source=manual)" - }, - "model_input": { - "name": "Model Input", - "tooltip": "Input from model loader node (only effective when model_source=auto)" - }, - "seed": { - "name": "Seed", - "tooltip": "Seed value (manual input or connect external node)" - }, - "positive_prompt": { - "name": "Positive Prompt", - "tooltip": "Positive prompt (manual input or connect external node)" - }, - "negative_prompt": { - "name": "Negative Prompt", - "tooltip": "Negative prompt (manual input or connect external node)" - }, - "custom_subfolder": { - "name": "Custom Subfolder", - "tooltip": "Custom subfolder name" - }, - "model_short_name": { - "name": "Model Short Name", - "tooltip": "Use model short name (remove extension)" - }, - "prompt_max_length": { - "name": "Prompt Max Length", - "tooltip": "Maximum length for prompt folder name" - }, - "sanitize_names": { - "name": "Sanitize Names", - "tooltip": "Clean illegal characters in filenames" - } - }, - "outputs": { - "0": { - "name": "Images", - "tooltip": "Passed image data" - }, - "1": { - "name": "Folder Path", - "tooltip": "Generated folder path" - }, - "2": { - "name": "Metadata JSON", - "tooltip": "Extracted metadata JSON" - } - } - }, - "SmartImageSaver": { - "display_name": "Smart Image Saver", - "description": "Smart image saver - handles image saving, format conversion and compression", - "inputs": { - "images": { - "name": "Images", - "tooltip": "Images to save (usually from SmartFolderManager)" - }, - "folder_path": { - "name": "Folder Path", - "tooltip": "Save path (from SmartFolderManager)" - }, - "metadata_json": { - "name": "Metadata JSON", - "tooltip": "Metadata JSON (from SmartFolderManager)" - }, - "filename_prefix": { - "name": "Filename Prefix", - "tooltip": "Filename prefix" - }, - "file_format": { - "name": "File Format", - "tooltip": "Image format", - "options": { - "png": "PNG Format", - "jpeg": "JPEG Format", - "webp": "WebP Format", - "bmp": "BMP Format", - "tiff": "TIFF Format" - } - }, - "preview_mode": { - "name": "Preview Mode", - "tooltip": "Save mode", - "options": { - "save_and_preview": "Save and Preview", - "preview_only": "Preview Only", - "save_only": "Save Only" - } - }, - "add_timestamp": { - "name": "Add Timestamp", - "tooltip": "Add timestamp to filename" - }, - "add_counter": { - "name": "Add Counter", - "tooltip": "Add counter (for multiple images)" - }, - "counter_start": { - "name": "Counter Start", - "tooltip": "Counter start value" - }, - "counter_padding": { - "name": "Counter Padding", - "tooltip": "Counter padding (zero-fill)" - }, - "jpeg_quality": { - "name": "JPEG Quality", - "tooltip": "JPEG quality (1-100)" - }, - "webp_quality": { - "name": "WebP Quality", - "tooltip": "WebP quality (1-100)" - }, - "webp_lossless": { - "name": "WebP Lossless", - "tooltip": "WebP lossless compression" - }, - "png_compression": { - "name": "PNG Compression", - "tooltip": "PNG compression level (0-9)" - }, - "embed_metadata": { - "name": "Embed Metadata", - "tooltip": "Embed metadata into image" - }, - "embed_workflow": { - "name": "Embed Workflow", - "tooltip": "Embed workflow information" - }, - "overwrite_existing": { - "name": "Overwrite Existing", - "tooltip": "Overwrite existing files" - }, - "create_backup": { - "name": "Create Backup", - "tooltip": "Create backup for overwritten files" - }, - "optimize_size": { - "name": "Optimize Size", - "tooltip": "Optimize file size" - } - }, - "outputs": { - "0": { - "name": "Images", - "tooltip": "Saved images" - } + "images": {"name": "Images"}, + "root_mode": {"name": "Save Location"}, + "custom_root": {"name": "Custom Path"}, + "folder_template": {"name": "Subfolder Rule"}, + "filename_template": {"name": "Filename Rule"}, + "file_format": {"name": "Format"}, + "collision_mode": {"name": "Name Collision"}, + "save_mode": {"name": "Save Mode"}, + "manual_model": {"name": "Model Source"}, + "embed_workflow": {"name": "Embed Workflow"}, + "counter_digits": {"name": "Counter Digits"}, + "png_compression": {"name": "PNG Compression"} } } -} \ No newline at end of file +} diff --git a/locales/zh/main.json b/locales/zh/main.json index ce64186..140baa8 100644 --- a/locales/zh/main.json +++ b/locales/zh/main.json @@ -1,4 +1,3 @@ { - "SmartFolderManager": "智能文件夹管理", - "SmartImageSaver": "智能图片保存" + "SmartSaveImage": "智能保存图片" } diff --git a/locales/zh/nodeDefs.json b/locales/zh/nodeDefs.json index 85d8493..a824015 100644 --- a/locales/zh/nodeDefs.json +++ b/locales/zh/nodeDefs.json @@ -1,206 +1,20 @@ { - "SmartFolderManager": { - "display_name": "智能文件夹管理", - "description": "智能文件夹管理器 - 负责管理文件夹结构和元数据收集", + "SmartSaveImage": { + "display_name": "智能保存图片", + "description": "预览最终路径，并用模板组织目录和文件名", "inputs": { - "images": { - "name": "图像", - "tooltip": "输入图片，用于提取尺寸信息并传递给保存节点" - }, - "base_folder": { - "name": "基础文件夹", - "tooltip": "基础文件夹路径，可以是相对路径或绝对路径" - }, - "create_subfolders": { - "name": "创建子文件夹", - "tooltip": "是否自动创建子文件夹" - }, - "enable_date_folder": { - "name": "启用日期文件夹", - "tooltip": "是否创建日期文件夹" - }, - "enable_model_folder": { - "name": "启用模型文件夹", - "tooltip": "是否创建模型文件夹" - }, - "enable_seed_folder": { - "name": "启用种子文件夹", - "tooltip": "是否创建种子文件夹" - }, - "enable_prompt_folder": { - "name": "启用提示词文件夹", - "tooltip": "是否创建提示词文件夹" - }, - "enable_custom_folder": { - "name": "启用自定义文件夹", - "tooltip": "是否使用自定义文件夹" - }, - "date_format": { - "name": "日期格式", - "tooltip": "日期格式: yyyy=年 MM=月 dd=日 hh=时 mm=分 ss=秒" - }, - "include_time": { - "name": "包含时间", - "tooltip": "是否在日期中包含时间" - }, - "model_source": { - "name": "模型来源", - "tooltip": "模型信息来源：auto=从元数据/外部输入获取，manual=手动选择", - "options": { - "auto": "自动获取", - "manual": "手动选择" - } - }, - "manual_model_name": { - "name": "手动模型名", - "tooltip": "手动选择模型（仅在model_source=manual时生效）" - }, - "model_input": { - "name": "模型输入", - "tooltip": "从模型加载器节点输入（仅在model_source=auto时生效）" - }, - "seed": { - "name": "随机种子", - "tooltip": "种子值（手动输入或连接外部节点）" - }, - "positive_prompt": { - "name": "正向提示词", - "tooltip": "正向提示词（手动输入或连接外部节点）" - }, - "negative_prompt": { - "name": "负向提示词", - "tooltip": "负向提示词（手动输入或连接外部节点）" - }, - "custom_subfolder": { - "name": "自定义子文件夹", - "tooltip": "自定义子文件夹名称" - }, - "model_short_name": { - "name": "模型短名称", - "tooltip": "使用模型短名称（去除扩展名）" - }, - "prompt_max_length": { - "name": "提示词最大长度", - "tooltip": "提示词文件夹名最大长度" - }, - "sanitize_names": { - "name": "清理文件名", - "tooltip": "清理文件名中的非法字符" - } - }, - "outputs": { - "0": { - "name": "图像", - "tooltip": "传递的图片数据" - }, - "1": { - "name": "文件夹路径", - "tooltip": "生成的文件夹路径" - }, - "2": { - "name": "元数据JSON", - "tooltip": "提取的元数据JSON" - } - } - }, - "SmartImageSaver": { - "display_name": "智能图片保存", - "description": "智能图片保存器 - 负责图片保存、格式转换和压缩", - "inputs": { - "images": { - "name": "图像", - "tooltip": "要保存的图片（通常来自SmartFolderManager）" - }, - "folder_path": { - "name": "文件夹路径", - "tooltip": "保存路径（来自SmartFolderManager）" - }, - "metadata_json": { - "name": "元数据JSON", - "tooltip": "元数据JSON（来自SmartFolderManager）" - }, - "filename_prefix": { - "name": "文件名前缀", - "tooltip": "文件名前缀" - }, - "file_format": { - "name": "文件格式", - "tooltip": "图片格式", - "options": { - "png": "PNG格式", - "jpeg": "JPEG格式", - "webp": "WebP格式", - "bmp": "BMP格式", - "tiff": "TIFF格式" - } - }, - "preview_mode": { - "name": "预览模式", - "tooltip": "保存模式", - "options": { - "save_and_preview": "保存并预览", - "preview_only": "仅预览", - "save_only": "仅保存" - } - }, - "add_timestamp": { - "name": "添加时间戳", - "tooltip": "在文件名中添加时间戳" - }, - "add_counter": { - "name": "添加计数器", - "tooltip": "添加计数器（多张图片时）" - }, - "counter_start": { - "name": "计数器起始值", - "tooltip": "计数器起始值" - }, - "counter_padding": { - "name": "计数器填充", - "tooltip": "计数器位数（补零）" - }, - "jpeg_quality": { - "name": "JPEG质量", - "tooltip": "JPEG质量（1-100）" - }, - "webp_quality": { - "name": "WebP质量", - "tooltip": "WebP质量（1-100）" - }, - "webp_lossless": { - "name": "WebP无损", - "tooltip": "WebP无损压缩" - }, - "png_compression": { - "name": "PNG压缩", - "tooltip": "PNG压缩级别（0-9）" - }, - "embed_metadata": { - "name": "嵌入元数据", - "tooltip": "嵌入元数据到图片" - }, - "embed_workflow": { - "name": "嵌入工作流", - "tooltip": "嵌入工作流信息" - }, - "overwrite_existing": { - "name": "覆盖已存在文件", - "tooltip": "覆盖已存在的文件" - }, - "create_backup": { - "name": "创建备份", - "tooltip": "为覆盖的文件创建备份" - }, - "optimize_size": { - "name": "优化文件大小", - "tooltip": "优化文件大小" - } - }, - "outputs": { - "0": { - "name": "图像", - "tooltip": "保存的图片" - } + "images": {"name": "图像", "tooltip": "需要保存的图片或图片批次"}, + "root_mode": {"name": "保存位置"}, + "custom_root": {"name": "自定义路径"}, + "folder_template": {"name": "子目录规则"}, + "filename_template": {"name": "文件名规则"}, + "file_format": {"name": "格式"}, + "collision_mode": {"name": "同名冲突"}, + "save_mode": {"name": "保存模式"}, + "manual_model": {"name": "模型来源"}, + "embed_workflow": {"name": "嵌入工作流"}, + "counter_digits": {"name": "序号位数"}, + "png_compression": {"name": "PNG 压缩"} } } -} \ No newline at end of file +} diff --git a/pyproject.toml b/pyproject.toml index efdb854..e9f8dbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta" [project] name = "SmartSaveImage" -version = "0.0.1" -description = "A node for easy save" +version = "2.0.0" +description = "A visual, template-driven image saver for ComfyUI" authors = [ {name = "kj", email = "2990346238@qq.com"} ] diff --git a/requirements.txt b/requirements.txt index ffbe029..2731c30 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,2 @@ Pillow>=8.0.0 -piexif>=1.1.3 numpy>=1.19.0 diff --git a/src/SmartSaveImage/__init__.py b/src/SmartSaveImage/__init__.py index b132cea..0fc8bb9 100644 --- a/src/SmartSaveImage/__init__.py +++ b/src/SmartSaveImage/__init__.py @@ -2,6 +2,14 @@ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +# 注册实时预览路由（在无 server 的测试环境中会静默跳过） +try: + from .server_routes import register_routes + + register_routes() +except Exception: # pragma: no cover + pass + __all__ = [ "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", diff --git a/src/SmartSaveImage/core/__init__.py b/src/SmartSaveImage/core/__init__.py deleted file mode 100644 index f179e41..0000000 --- a/src/SmartSaveImage/core/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -"""SmartSaveImage 核心模块""" - -from .metadata import MetadataExtractor, MetadataBuilder -from .path_utils import PathManager -from .image_utils import ImageProcessor - -__all__ = [ - "MetadataExtractor", - "MetadataBuilder", - "PathManager", - "ImageProcessor", -] diff --git a/src/SmartSaveImage/core/image_utils.py b/src/SmartSaveImage/core/image_utils.py deleted file mode 100644 index f61e661..0000000 --- a/src/SmartSaveImage/core/image_utils.py +++ /dev/null @@ -1,168 +0,0 @@ -"""图片处理工具模块""" - -import os -import json -import numpy as np -from PIL import Image, PngImagePlugin -import piexif -from typing import Dict, Any, Optional, Tuple - - -class ImageProcessor: - """图片处理器""" - - def __init__(self): - self.supported_formats = { - "png": {"extension": ".png", "pil_format": "PNG"}, - "jpeg": {"extension": ".jpg", "pil_format": "JPEG"}, - "webp": {"extension": ".webp", "pil_format": "WEBP"}, - "bmp": {"extension": ".bmp", "pil_format": "BMP"}, - "tiff": {"extension": ".tiff", "pil_format": "TIFF"}, - } - - def tensor_to_pil(self, tensor_image) -> Image.Image: - """将tensor图片转换为PIL图片""" - # 处理批次维度 - if len(tensor_image.shape) == 4: - tensor_image = tensor_image.squeeze(0) - - # 转换为numpy数组并调整范围到0-255 - image_np = (tensor_image.cpu().numpy() * 255).astype(np.uint8) - - # 转换为PIL图片 - return Image.fromarray(image_np) - - def prepare_save_kwargs(self, file_format: str, quality_settings: Dict[str, Any]) -> Dict[str, Any]: - """准备保存参数""" - format_info = self.supported_formats.get(file_format, self.supported_formats["png"]) - save_kwargs = {"format": format_info["pil_format"]} - - if file_format == "png": - save_kwargs["compress_level"] = quality_settings.get("png_compression", 6) - save_kwargs["optimize"] = quality_settings.get("optimize_size", False) - - elif file_format == "jpeg": - save_kwargs["quality"] = quality_settings.get("jpeg_quality", 95) - save_kwargs["optimize"] = quality_settings.get("optimize_size", False) - - elif file_format == "webp": - if quality_settings.get("webp_lossless", False): - save_kwargs["lossless"] = True - else: - save_kwargs["quality"] = quality_settings.get("webp_quality", 90) - save_kwargs["method"] = 6 # 最佳压缩方法 - - return save_kwargs - - def add_png_metadata(self, pil_image: Image.Image, metadata_text: Optional[str], - workflow_data: Optional[Dict], save_kwargs: Dict) -> Dict: - """为PNG格式添加元数据""" - if metadata_text or workflow_data: - pnginfo = PngImagePlugin.PngInfo() - - if metadata_text: - pnginfo.add_text("parameters", metadata_text) - - if workflow_data: - pnginfo.add_text("workflow", json.dumps(workflow_data)) - - save_kwargs["pnginfo"] = pnginfo - - return save_kwargs - - def add_exif_metadata(self, metadata_text: Optional[str], - workflow_data: Optional[Dict]) -> Optional[bytes]: - """创建EXIF元数据（用于JPEG和WebP）""" - if not metadata_text and not workflow_data: - return None - - try: - exif_dict = {} - - if metadata_text: - # 将元数据添加到UserComment字段 - exif_dict["Exif"] = { - piexif.ExifIFD.UserComment: b"UNICODE\0" + metadata_text.encode("utf-16be") - } - - if workflow_data: - # 将工作流数据添加到ImageDescription字段 - workflow_json = json.dumps(workflow_data) - exif_dict["0th"] = { - piexif.ImageIFD.ImageDescription: f"Workflow:{workflow_json}" - } - - return piexif.dump(exif_dict) - - except Exception as e: - print(f"[ImageProcessor] EXIF元数据创建失败: {e}") - return None - - def save_image(self, tensor_image, filepath: str, file_format: str, - quality_settings: Dict[str, Any], metadata_text: Optional[str] = None, - embed_workflow: bool = False, workflow_data: Optional[Dict] = None) -> bool: - """保存图片并嵌入元数据""" - try: - # 转换为PIL图片 - pil_image = self.tensor_to_pil(tensor_image) - - # 准备保存参数 - save_kwargs = self.prepare_save_kwargs(file_format, quality_settings) - - # 根据格式添加元数据 - if file_format == "png": - # PNG使用PngInfo - save_kwargs = self.add_png_metadata( - pil_image, metadata_text, - workflow_data if embed_workflow else None, - save_kwargs - ) - - elif file_format in ["jpeg", "webp"]: - # JPEG和WebP使用EXIF - exif_bytes = self.add_exif_metadata( - metadata_text, - workflow_data if embed_workflow else None - ) - if exif_bytes: - save_kwargs["exif"] = exif_bytes - - # 保存图片 - pil_image.save(filepath, **save_kwargs) - return True - - except Exception as e: - print(f"[ImageProcessor] 保存图片失败 {filepath}: {e}") - return False - - def create_backup(self, filepath: str) -> bool: - """为现有文件创建备份""" - if not os.path.exists(filepath): - return True - - try: - backup_path = filepath + ".backup" - - # 如果备份已存在，先删除 - if os.path.exists(backup_path): - os.remove(backup_path) - - # 重命名原文件为备份 - os.rename(filepath, backup_path) - print(f"[ImageProcessor] 创建备份: {backup_path}") - return True - - except Exception as e: - print(f"[ImageProcessor] 创建备份失败: {e}") - return False - - def get_image_info(self, tensor_image) -> Tuple[int, int]: - """获取图片尺寸信息""" - if len(tensor_image.shape) == 4: - # 批次格式: (batch, height, width, channels) - return tensor_image.shape[2], tensor_image.shape[1] # width, height - elif len(tensor_image.shape) == 3: - # 单张格式: (height, width, channels) - return tensor_image.shape[1], tensor_image.shape[0] # width, height - else: - return 0, 0 diff --git a/src/SmartSaveImage/core/metadata.py b/src/SmartSaveImage/core/metadata.py deleted file mode 100644 index eea3e37..0000000 --- a/src/SmartSaveImage/core/metadata.py +++ /dev/null @@ -1,201 +0,0 @@ -"""元数据提取和构建模块""" - -import json -import os -from datetime import datetime -from typing import Dict, Any, Optional - - -class MetadataExtractor: - """从ComfyUI工作流中提取元数据""" - - def __init__(self): - self.supported_checkpoint_nodes = [ - "CheckpointLoaderSimple", - "CheckpointLoader", - "CheckpointLoaderV2" - ] - self.supported_sampler_nodes = [ - "KSampler", - "KSamplerAdvanced", - "SamplerCustom", - "SamplerCustomAdvanced", - # 添加更多可能的采样器节点类型 - ] - self.supported_text_nodes = [ - "CLIPTextEncode", - "CLIPTextEncodeSDXL", - "T5TextEncode" - ] - - def extract_from_workflow(self, extra_pnginfo: Optional[Dict]) -> Dict[str, Any]: - """从工作流信息中提取元数据""" - metadata = { - "model": None, - "seed": None, - "steps": None, - "cfg": None, - "sampler": None, - "scheduler": None, - "positive_prompt": None, - "negative_prompt": None, - "width": None, - "height": None, - } - - if not isinstance(extra_pnginfo, dict): - return metadata - - workflow = extra_pnginfo.get("workflow", {}) - if not isinstance(workflow, dict): - return metadata - - nodes_data = workflow.get("nodes", []) - text_prompts = [] # 收集所有文本提示 - - for node in nodes_data: - if not isinstance(node, dict): - continue - - class_type = node.get("class_type", "") - inputs = node.get("inputs", {}) - - # 提取模型信息 - if class_type in self.supported_checkpoint_nodes: - if "ckpt_name" in inputs and not metadata["model"]: - metadata["model"] = inputs["ckpt_name"] - - # 提取采样器信息 - elif class_type in self.supported_sampler_nodes: - if "seed" in inputs and metadata["seed"] is None: - metadata["seed"] = inputs["seed"] - if "steps" in inputs and metadata["steps"] is None: - metadata["steps"] = inputs["steps"] - if "cfg" in inputs and metadata["cfg"] is None: - metadata["cfg"] = inputs["cfg"] - if "sampler_name" in inputs and not metadata["sampler"]: - metadata["sampler"] = inputs["sampler_name"] - if "scheduler" in inputs and not metadata["scheduler"]: - metadata["scheduler"] = inputs["scheduler"] - - # 提取文本提示 - elif class_type in self.supported_text_nodes: - text = inputs.get("text", "") - if text and isinstance(text, str): - text_prompts.append(text.strip()) - - # 分配正负提示词（通常第一个是正向，第二个是负向） - if text_prompts: - metadata["positive_prompt"] = text_prompts[0] - if len(text_prompts) > 1: - metadata["negative_prompt"] = text_prompts[1] - - return metadata - - def extract_from_prompt(self, prompt: Optional[Dict]) -> Dict[str, Any]: - """从prompt参数中提取元数据（备用方案）""" - metadata = { - "model": None, - "seed": None, - "steps": None, - "cfg": None, - "sampler": None, - "scheduler": None, - } - - if not isinstance(prompt, dict): - return metadata - - for node_id, node_data in prompt.items(): - if not isinstance(node_data, dict): - continue - - class_type = node_data.get("class_type", "") - inputs = node_data.get("inputs", {}) - - if class_type in self.supported_checkpoint_nodes and not metadata["model"]: - metadata["model"] = inputs.get("ckpt_name") - elif class_type in self.supported_sampler_nodes: - if metadata["seed"] is None: - metadata["seed"] = inputs.get("seed") - if metadata["steps"] is None: - metadata["steps"] = inputs.get("steps") - if metadata["cfg"] is None: - metadata["cfg"] = inputs.get("cfg") - if not metadata["sampler"]: - metadata["sampler"] = inputs.get("sampler_name") - if not metadata["scheduler"]: - metadata["scheduler"] = inputs.get("scheduler") - - return metadata - - -class MetadataBuilder: - """构建用于保存的元数据""" - - def build_parameters_text(self, metadata: Dict[str, Any]) -> str: - """构建参数文本（用于嵌入图片）""" - parts = [] - - # 正向提示词 - positive = metadata.get("positive_prompt") - if positive: - parts.append(str(positive)) - - # 负向提示词 - negative = metadata.get("negative_prompt") - if negative: - parts.append(f"Negative prompt: {negative}") - - # 技术参数 - params = [] - - steps = metadata.get("steps") - if steps is not None: - params.append(f"Steps: {steps}") - - sampler = metadata.get("sampler") - scheduler = metadata.get("scheduler") - if sampler: - if scheduler and scheduler != "normal": - params.append(f"Sampler: {sampler} {scheduler}") - else: - params.append(f"Sampler: {sampler}") - - cfg = metadata.get("cfg") - if cfg is not None: - params.append(f"CFG scale: {cfg}") - - seed = metadata.get("seed") - if seed is not None: - params.append(f"Seed: {seed}") - - width = metadata.get("width") - height = metadata.get("height") - if width and height: - params.append(f"Size: {width}x{height}") - - model = metadata.get("model") - if model: - model_name = os.path.splitext(os.path.basename(model))[0] - # 这里可以添加模型哈希计算，但为了性能考虑暂时省略 - params.append(f"Model: {model_name}") - - if params: - parts.append(", ".join(params)) - - return "\n".join(parts) - - def build_metadata_json(self, folder_path: str, structure_mode: str, - workflow_metadata: Dict, user_inputs: Dict) -> str: - """构建完整的元数据JSON""" - metadata = { - "folder_path": folder_path, - "structure_mode": structure_mode, - "timestamp": datetime.now().isoformat(), - "workflow_metadata": workflow_metadata, - "user_inputs": user_inputs, - "version": "1.0" - } - - return json.dumps(metadata, ensure_ascii=False, indent=2) diff --git a/src/SmartSaveImage/core/path_utils.py b/src/SmartSaveImage/core/path_utils.py deleted file mode 100644 index 82c28c1..0000000 --- a/src/SmartSaveImage/core/path_utils.py +++ /dev/null @@ -1,230 +0,0 @@ -"""路径管理和文件名处理工具""" - -import os -import re -from datetime import datetime -from typing import List, Dict, Any, Optional - - -class PathManager: - """路径管理器""" - - def __init__(self): - # Windows文件名非法字符 - self.illegal_chars = r'[<>:"/\\|?*]' - # 日期格式替换表 - self.date_formats = { - "yyyy": lambda dt: f"{dt.year:04d}", - "yy": lambda dt: f"{dt.year % 100:02d}", - "MM": lambda dt: f"{dt.month:02d}", - "dd": lambda dt: f"{dt.day:02d}", - "hh": lambda dt: f"{dt.hour:02d}", - "mm": lambda dt: f"{dt.minute:02d}", - "ss": lambda dt: f"{dt.second:02d}", - } - - def sanitize_filename(self, name: str, max_length: int = 100) -> str: - """安全地清理文件名，移除非法字符和路径遍历攻击""" - if not name: - return "untitled" - - name = str(name) - - # 使用os.path.basename确保只获取文件名部分，防止路径遍历 - name = os.path.basename(name) - - # 如果basename后为空（比如输入是路径分隔符），使用默认名称 - if not name or name in ('.', '..'): - return "untitled" - - # 移除或替换非法字符 - name = re.sub(self.illegal_chars, '_', name) - - # 处理空格和特殊字符 - name = re.sub(r'\s+', '_', name) # 多个空格替换为下划线 - name = re.sub(r'[^\w\-_.]', '_', name) # 只保留字母数字下划线破折号点 - name = re.sub(r'_+', '_', name) # 多个下划线合并为一个 - - # 移除首尾的下划线和点 - name = name.strip('_.') - - # 防止Windows保留名称 - reserved_names = [ - 'CON', 'PRN', 'AUX', 'NUL', - 'COM1', 'COM2', 'COM3', 'COM4', 'COM5', 'COM6', 'COM7', 'COM8', 'COM9', - 'LPT1', 'LPT2', 'LPT3', 'LPT4', 'LPT5', 'LPT6', 'LPT7', 'LPT8', 'LPT9' - ] - name_upper = name.upper() - if name_upper in reserved_names or name_upper.split('.')[0] in reserved_names: - name = f"_{name}" - - # 限制长度 - if len(name) > max_length: - name = name[:max_length].rstrip('_.') - - return name if name else "untitled" - - def format_date(self, date_format: str, include_time: bool = False) -> str: - """格式化日期字符串""" - now = datetime.now() - - # 如果需要包含时间但格式中没有时间部分，自动添加 - if include_time and not any(t in date_format for t in ["hh", "mm", "ss"]): - date_format += "_hh-mm-ss" - - # 替换日期格式标记 - result = date_format - for pattern, formatter in self.date_formats.items(): - result = result.replace(pattern, formatter(now)) - - return result - - def build_folder_structure(self, base_path: str, structure_mode: str, - metadata: Dict[str, Any], user_inputs: Dict[str, Any]) -> List[str]: - """构建文件夹结构路径段""" - path_segments = [] - - if structure_mode == "date": - date_str = self.format_date( - user_inputs.get("date_format", "yyyy-MM-dd"), - user_inputs.get("include_time", False) - ) - path_segments.append(date_str) - - elif structure_mode == "model": - model = user_inputs.get("model_name") or metadata.get("model") - if model: - if user_inputs.get("model_short_name", True): - model = os.path.splitext(os.path.basename(model))[0] - path_segments.append(self.sanitize_filename(model)) - - elif structure_mode == "seed": - seed = user_inputs.get("seed_value") - if seed is None: - seed = metadata.get("seed") - if seed is not None: - path_segments.append(f"seed_{seed}") - - elif structure_mode == "prompt": - prompt = user_inputs.get("prompt_text") or metadata.get("positive_prompt") - if prompt: - max_len = user_inputs.get("prompt_max_length", 50) - prompt_clean = str(prompt).replace("\n", " ").strip()[:max_len] - path_segments.append(self.sanitize_filename(prompt_clean)) - - elif structure_mode == "custom": - custom_path = user_inputs.get("custom_path", "") - if custom_path: - # 支持变量替换 - variables = { - "{date}": self.format_date(user_inputs.get("date_format", "yyyy-MM-dd")), - "{model}": self.sanitize_filename( - user_inputs.get("model_name") or metadata.get("model", "model") - ), - "{seed}": str(user_inputs.get("seed_value") or metadata.get("seed", 0)), - "{prompt}": self.sanitize_filename( - (user_inputs.get("prompt_text") or metadata.get("positive_prompt", ""))[:50] - ), - } - - for var, value in variables.items(): - custom_path = custom_path.replace(var, value) - - # 分割路径并清理 - segments = [s.strip() for s in custom_path.split("/") if s.strip()] - path_segments.extend([self.sanitize_filename(s) for s in segments]) - - elif structure_mode == "auto": - # 自动模式：日期/模型/种子的组合 - date_str = self.format_date("yyyy-MM-dd") - path_segments.append(date_str) - - # 添加模型 - model = user_inputs.get("model_name") or metadata.get("model") - if model: - if user_inputs.get("model_short_name", True): - model = os.path.splitext(os.path.basename(model))[0] - path_segments.append(self.sanitize_filename(model)) - - # 添加种子 - seed = user_inputs.get("seed_value") - if seed is None: - seed = metadata.get("seed") - if seed is not None: - path_segments.append(f"seed_{seed}") - - # 限制文件夹深度 - max_depth = user_inputs.get("max_folder_depth", 5) - path_segments = path_segments[:max_depth] - - return [seg for seg in path_segments if seg] - - def resolve_base_path(self, base_folder: str, default_output_dir: str) -> str: - """安全地解析基础路径""" - from ..utils.validators import InputValidator - - if not base_folder or base_folder.lower() in ["", "output", "."]: - return default_output_dir - - # 使用安全的路径连接方法 - return InputValidator.secure_path_join(base_folder, "") - - def build_full_path(self, base_folder: str, structure_mode: str, - metadata: Dict[str, Any], user_inputs: Dict[str, Any], - default_output_dir: str) -> str: - """构建完整的文件夹路径""" - # 解析基础路径 - base_path = self.resolve_base_path(base_folder, default_output_dir) - - # 构建子文件夹结构 - path_segments = self.build_folder_structure(base_path, structure_mode, metadata, user_inputs) - - if path_segments: - final_path = os.path.join(base_path, *path_segments) - else: - final_path = base_path - from ..utils.validators import InputValidator - return InputValidator.secure_path_join(final_path, "") - - def ensure_directory_exists(self, path: str) -> bool: - """确保目录存在，如果不存在则创建""" - try: - os.makedirs(path, exist_ok=True) - return True - except Exception as e: - print(f"[PathManager] 创建目录失败: {path}, 错误: {e}") - return False - - def generate_unique_filename(self, directory: str, base_filename: str, - extension: str, overwrite: bool = False) -> str: - filename = f"{base_filename}{extension}" - filepath = os.path.join(directory, filename) - if not os.path.exists(filepath) or overwrite: - return filename - m = re.match(r"^(.*)_(\d+)$", base_filename) - if m: - root = m.group(1) - num_str = m.group(2) - padding = len(num_str) - try: - num = int(num_str) - except ValueError: - num = 1 - while True: - num += 1 - new_filename = f"{root}_{str(num).zfill(padding)}{extension}" - if not os.path.exists(os.path.join(directory, new_filename)): - return new_filename - if num > 9999999: - break - else: - counter = 1 - while True: - new_filename = f"{base_filename}_{counter:03d}{extension}" - if not os.path.exists(os.path.join(directory, new_filename)): - return new_filename - counter += 1 - if counter > 9999999: - break - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:17] - return f"{base_filename}_{timestamp}{extension}" diff --git a/src/SmartSaveImage/nodes.py b/src/SmartSaveImage/nodes.py deleted file mode 100644 index fab5cff..0000000 --- a/src/SmartSaveImage/nodes.py +++ /dev/null @@ -1,118 +0,0 @@ -from inspect import cleandoc -class Example: - """ - A example node - - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. - IS_CHANGED: - optional method to control when the node is re executed. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tulple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tulple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. - - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Secound value is a config for type "INT", "STRING" or "FLOAT". - """ - return { - "required": { - "image": ("Image", { "tooltip": "This is an image"}), - "int_field": ("INT", { - "default": 0, - "min": 0, #Minimum value - "max": 4096, #Maximum value - "step": 64, #Slider's step - "display": "number" # Cosmetic only: display as "number" or "slider" - }), - "float_field": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 10.0, - "step": 0.01, - "round": 0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. - "display": "number"}), - "print_to_screen": (["enable", "disable"],), - "string_field": ("STRING", { - "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node - "default": "Hello World!" - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - #RETURN_NAMES = ("image_output_name",) - DESCRIPTION = cleandoc(__doc__) - FUNCTION = "test" - - #OUTPUT_NODE = False - #OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node - - CATEGORY = "Example" - - def test(self, image, string_field, int_field, float_field, print_to_screen): - if print_to_screen == "enable": - print(f"""Your input contains: - string_field aka input text: {string_field} - int_field: {int_field} - float_field: {float_field} - """) - #do some processing on the image, in this example I just invert it - image = 1.0 - image - return (image,) - - """ - The node will always be re executed if any of the inputs change but - this method can be used to force the node to execute again even when the inputs don't change. - You can make this node return a number or a string. This value will be compared to the one returned the last time the node was - executed, if it is different the node will be executed again. - This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash - changes between executions the LoadImage node is executed again. - """ - #@classmethod - #def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen): - # return "" - - -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique -NODE_CLASS_MAPPINGS = { - "Example": Example -} - -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = { - "Example": "Example Node" -} diff --git a/src/SmartSaveImage/nodes/__init__.py b/src/SmartSaveImage/nodes/__init__.py index 9e8fa0e..f0d8a76 100644 --- a/src/SmartSaveImage/nodes/__init__.py +++ b/src/SmartSaveImage/nodes/__init__.py @@ -1,22 +1,18 @@ """SmartSaveImage 节点模块""" -from .folder_manager import SmartFolderManager -from .image_saver import SmartImageSaver +from .smart_save import SmartSaveImage -# 导出节点映射 +# 节点 ID 保持为 SmartSaveImage，旧工作流可无缝复用 NODE_CLASS_MAPPINGS = { - "SmartFolderManager": SmartFolderManager, - "SmartImageSaver": SmartImageSaver, + "SmartSaveImage": SmartSaveImage, } NODE_DISPLAY_NAME_MAPPINGS = { - "SmartFolderManager": "SmartFolderManager", - "SmartImageSaver": "SmartImageSaver", + "SmartSaveImage": "智能保存图片", } __all__ = [ - "SmartFolderManager", - "SmartImageSaver", + "SmartSaveImage", "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", ] diff --git a/src/SmartSaveImage/nodes/folder_manager.py b/src/SmartSaveImage/nodes/folder_manager.py deleted file mode 100644 index c2f1690..0000000 --- a/src/SmartSaveImage/nodes/folder_manager.py +++ /dev/null @@ -1,412 +0,0 @@ -"""智能文件夹管理器节点""" - -import os -import folder_paths -from ..core import MetadataExtractor, MetadataBuilder, PathManager -from ..utils import InputValidator - - -class SmartFolderManager: - """智能文件夹管理器 - 负责管理文件夹结构和元数据收集""" - - CATEGORY = "SmartSave" - RETURN_TYPES = ("IMAGE", "STRING", "STRING") - RETURN_NAMES = ("images", "folder_path", "metadata_json") - FUNCTION = "generate_path" - - # 为翻译插件提供类名标识 - @classmethod - def get_comfy_class(cls): - return "SmartFolderManager" - - def __init__(self): - self.metadata_extractor = MetadataExtractor() - self.metadata_builder = MetadataBuilder() - self.path_manager = PathManager() - self.validator = InputValidator() - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "images": ("IMAGE", {"tooltip": "Input images for extracting size info and passing to save node"}), - "base_folder": ("STRING", { - "default": "output", - "multiline": False, - "tooltip": "Base folder path, can be relative or absolute path" - }), - "create_subfolders": ("BOOLEAN", { - "default": True, - "tooltip": "Whether to automatically create subfolders" - }), - }, - "optional": { - # 文件夹层级开关 - "enable_date_folder": ("BOOLEAN", { - "default": True, - "tooltip": "Whether to create date folder" - }), - "enable_model_folder": ("BOOLEAN", { - "default": True, - "tooltip": "Whether to create model folder" - }), - "enable_seed_folder": ("BOOLEAN", { - "default": False, - "tooltip": "Whether to create seed folder" - }), - "enable_prompt_folder": ("BOOLEAN", { - "default": False, - "tooltip": "Whether to create prompt folder" - }), - "enable_custom_folder": ("BOOLEAN", { - "default": False, - "tooltip": "Whether to use custom folder" - }), - - # 日期相关 - "date_format": ("STRING", { - "default": "yyyy-MM-dd", - "multiline": False, - "tooltip": "Date format: yyyy=year MM=month dd=day hh=hour mm=minute ss=second" - }), - "include_time": ("BOOLEAN", { - "default": False, - "tooltip": "Whether to include time in date" - }), - - # 模型信息来源选择 - "model_source": (["auto", "manual"], { - "default": "auto", - "tooltip": "Model info source: auto=from metadata/external input, manual=manual selection" - }), - "manual_model_name": (folder_paths.get_filename_list("checkpoints"), { - "tooltip": "Manual model selection (only effective when model_source=manual)" - }), - "model_input": ("MODEL", { - "tooltip": "Input from model loader node (only effective when model_source=auto)" - }), - - # 种子信息 - "seed": ("STRING", { - "default": "0", - "multiline": False, - "tooltip": "Seed value (manual input or connect external node)" - }), - - # 提示词信息 - "positive_prompt": ("STRING", { - "default": "", - "multiline": True, - "tooltip": "Positive prompt (manual input or connect external node)" - }), - "negative_prompt": ("STRING", { - "default": "", - "multiline": True, - "tooltip": "Negative prompt (manual input or connect external node)" - }), - - # 自定义路径 - "custom_subfolder": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Custom subfolder name" - }), - - # 显示选项 - "model_short_name": ("BOOLEAN", { - "default": True, - "tooltip": "Use model short name (remove extension)" - }), - "prompt_max_length": ("INT", { - "default": 50, - "min": 10, - "max": 200, - "tooltip": "Maximum length for prompt folder name" - }), - "sanitize_names": ("BOOLEAN", { - "default": True, - "tooltip": "Clean illegal characters in filenames" - }), - }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - }, - } - - def generate_path(self, images, base_folder, create_subfolders, - enable_date_folder=True, enable_model_folder=True, - enable_seed_folder=False, enable_prompt_folder=False, enable_custom_folder=False, - date_format="yyyy-MM-dd", include_time=False, - model_source="auto", manual_model_name=None, model_input=None, - seed="0", positive_prompt="", negative_prompt="", - custom_subfolder="", model_short_name=True, prompt_max_length=50, sanitize_names=True, - prompt=None, extra_pnginfo=None): - """生成文件夹路径和元数据""" - - # 验证输入 - if not self.validator.validate_folder_path(base_folder): - print(f"[SmartFolderManager] 无效的基础文件夹路径: {base_folder}") - base_folder = "output" - - if date_format and not self.validator.validate_date_format(date_format): - print(f"[SmartFolderManager] 无效的日期格式: {date_format}") - date_format = "yyyy-MM-dd" - - # 清理输入字符串 - positive_prompt = self.validator.sanitize_input_string(positive_prompt, prompt_max_length * 2) - negative_prompt = self.validator.sanitize_input_string(negative_prompt, prompt_max_length * 2) - custom_subfolder = self.validator.sanitize_input_string(custom_subfolder, 100) - - # 从图片中提取尺寸信息 - image_metadata = {} - if images is not None and len(images) > 0: - try: - # 获取第一张图片的尺寸 - if len(images.shape) == 4: # [batch, height, width, channels] - height, width = images.shape[1], images.shape[2] - elif len(images.shape) == 3: # [height, width, channels] - height, width = images.shape[0], images.shape[1] - else: - height, width = 0, 0 - - image_metadata["width"] = width - image_metadata["height"] = height - print(f"[SmartFolderManager] 从图片提取尺寸: {width}x{height}") - except Exception as e: - print(f"[SmartFolderManager] 提取图片尺寸失败: {e}") - image_metadata["width"] = 0 - image_metadata["height"] = 0 - - # 1. 根据用户选择构建元数据 - final_metadata = {} - - # 模型信息 - if model_source == "manual": - final_metadata["model"] = manual_model_name - else: # auto - # 优先从外部输入获取 - if model_input is not None: - external_model = self.extract_model_from_input(model_input) - if external_model: - final_metadata["model"] = external_model - else: - # 从工作流元数据获取 - workflow_metadata = self.metadata_extractor.extract_from_prompt(prompt) - final_metadata["model"] = workflow_metadata.get("model") - else: - # 从工作流元数据获取 - workflow_metadata = self.metadata_extractor.extract_from_prompt(prompt) - final_metadata["model"] = workflow_metadata.get("model") - - # 种子信息 - final_metadata["seed"] = seed - - # 提示词信息 - final_metadata["positive_prompt"] = positive_prompt if positive_prompt.strip() else None - final_metadata["negative_prompt"] = negative_prompt if negative_prompt.strip() else None - - # 尺寸信息（直接从图片获取） - final_metadata["width"] = image_metadata.get("width", 0) - final_metadata["height"] = image_metadata.get("height", 0) - - # 2. 构建文件夹路径 - path_segments = [] - base_path = self.path_manager.resolve_base_path(base_folder, folder_paths.get_output_directory()) - - # 按开关添加各层文件夹 - if enable_date_folder: - date_str = self.path_manager.format_date(date_format, include_time) - path_segments.append(date_str) - - if enable_model_folder and final_metadata.get("model"): - model_name = final_metadata["model"] - if model_short_name: - model_name = os.path.splitext(os.path.basename(model_name))[0] - if sanitize_names: - model_name = self.path_manager.sanitize_filename(model_name) - path_segments.append(model_name) - - if enable_seed_folder and final_metadata.get("seed") is not None: - seed_str = f"seed_{final_metadata['seed']}" - path_segments.append(seed_str) - - if enable_prompt_folder and final_metadata.get("positive_prompt"): - prompt_text = final_metadata["positive_prompt"] - prompt_clean = prompt_text.replace("\n", " ").strip()[:prompt_max_length] - if sanitize_names: - prompt_clean = self.path_manager.sanitize_filename(prompt_clean) - path_segments.append(prompt_clean) - - if enable_custom_folder and custom_subfolder: - if sanitize_names: - custom_subfolder = self.path_manager.sanitize_filename(custom_subfolder) - path_segments.append(custom_subfolder) - - if path_segments: - folder_path = os.path.join(base_path, *path_segments) - else: - folder_path = base_path - folder_path = InputValidator.secure_path_join(folder_path, "") - - # 创建文件夹 - if create_subfolders: - if not self.path_manager.ensure_directory_exists(folder_path): - print(f"[SmartFolderManager] 创建文件夹失败，使用默认输出目录") - folder_path = folder_paths.get_output_directory() - else: - print(f"[SmartFolderManager] 文件夹路径: {folder_path}") - - # 构建用户输入记录 - user_inputs = { - "enable_date_folder": enable_date_folder, - "enable_model_folder": enable_model_folder, - "enable_seed_folder": enable_seed_folder, - "enable_prompt_folder": enable_prompt_folder, - "enable_custom_folder": enable_custom_folder, - "date_format": date_format, - "include_time": include_time, - "model_source": model_source, - "manual_model_name": manual_model_name, - "seed": seed, - "positive_prompt": positive_prompt, - "negative_prompt": negative_prompt, - "custom_subfolder": custom_subfolder, - } - - # 构建元数据JSON - metadata_json = self.metadata_builder.build_metadata_json( - folder_path, "flexible", final_metadata, user_inputs - ) - - return (images, folder_path, metadata_json) - - def extract_model_from_input(self, model_input): - """从模型输入中提取模型名称（增强版）""" - if model_input is None: - return None - - try: - # 检查ModelPatcher的关键属性 - key_attrs = ['model_path', 'checkpoint_path', 'ckpt_path', 'model_file', 'filename'] - for attr in key_attrs: - if hasattr(model_input, attr): - value = getattr(model_input, attr) - if value: - return value - - # 尝试从ModelPatcher的特殊属性中提取信息 - # 检查model_options（可能包含模型信息） - if hasattr(model_input, 'model_options') and model_input.model_options: - if isinstance(model_input.model_options, dict): - for key in ['model_path', 'checkpoint_path', 'filename']: - if key in model_input.model_options: - value = model_input.model_options[key] - if value: - return value - - # 检查attachments（可能包含模型信息） - if hasattr(model_input, 'attachments') and model_input.attachments: - if isinstance(model_input.attachments, dict): - for key, value in model_input.attachments.items(): - if isinstance(value, str) and ('.safetensors' in value or '.ckpt' in value): - return value - - # 最后的尝试：检查parent属性 - if hasattr(model_input, 'parent') and model_input.parent: - parent_model = self.extract_model_from_input(model_input.parent) - if parent_model: - return parent_model - - # 字典格式 - if isinstance(model_input, dict): - return model_input.get('model_path') or model_input.get('checkpoint_path') - - # 尝试其他可能的属性名 - possible_attrs = ['checkpoint_path', 'ckpt_path', 'model_file', 'filename', 'path'] - for attr in possible_attrs: - if hasattr(model_input, attr): - value = getattr(model_input, attr) - if value: - return value - - # 如果是tuple或list，检查第一个元素 - if isinstance(model_input, (tuple, list)) and len(model_input) > 0: - return self.extract_model_from_input(model_input[0]) - - except Exception as e: - pass # 静默处理错误 - - return None - - def extract_from_external_inputs(self, model_input, conditioning_positive, conditioning_negative, latent_input): - """从外部输入节点提取元数据""" - metadata = {} - - # 从模型输入提取模型名称 - if model_input is not None: - try: - # 尝试多种方式从模型对象中提取信息 - model_name = None - - # 方法1: 检查是否有model_path属性 - if hasattr(model_input, 'model_path'): - model_name = model_input.model_path - - # 方法2: 检查model对象的属性 - elif hasattr(model_input, 'model'): - model_obj = model_input.model - if hasattr(model_obj, 'model_path'): - model_name = model_obj.model_path - elif hasattr(model_obj, 'checkpoint_path'): - model_name = model_obj.checkpoint_path - - # 方法3: 检查是否是字典格式 - elif isinstance(model_input, dict): - model_name = model_input.get('model_path') or model_input.get('checkpoint_path') - - if model_name: - metadata["model"] = model_name - print(f"[SmartFolderManager] 从外部输入提取模型: {model_name}") - else: - print(f"[SmartFolderManager] 模型输入已连接，但无法提取模型名称") - - except Exception as e: - print(f"[SmartFolderManager] 从模型输入提取信息失败: {e}") - - # 从conditioning提取提示词 - # 注意：ComfyUI的conditioning对象通常不直接包含原始文本 - # 但我们可以尝试一些方法 - if conditioning_positive is not None: - try: - # conditioning通常是一个包含编码后数据的复杂结构 - # 我们标记有外部conditioning输入，但文本提取仍依赖工作流 - metadata["has_positive_conditioning"] = True - print(f"[SmartFolderManager] 检测到正向条件输入") - except Exception as e: - print(f"[SmartFolderManager] 处理正向条件输入失败: {e}") - - if conditioning_negative is not None: - try: - metadata["has_negative_conditioning"] = True - print(f"[SmartFolderManager] 检测到负向条件输入") - except Exception as e: - print(f"[SmartFolderManager] 处理负向条件输入失败: {e}") - - # 从latent提取尺寸信息 - if latent_input is not None: - try: - if isinstance(latent_input, dict) and "samples" in latent_input: - samples = latent_input["samples"] - if hasattr(samples, 'shape') and len(samples.shape) >= 3: - # latent通常是 [batch, channels, height, width] - # 需要乘以8因为latent空间是1/8分辨率 - height = samples.shape[-2] * 8 - width = samples.shape[-1] * 8 - metadata["width"] = width - metadata["height"] = height - print(f"[SmartFolderManager] 从latent提取尺寸: {width}x{height}") - except Exception as e: - print(f"[SmartFolderManager] 从latent提取尺寸失败: {e}") - - return metadata diff --git a/src/SmartSaveImage/nodes/image_saver.py b/src/SmartSaveImage/nodes/image_saver.py deleted file mode 100644 index 1a8877e..0000000 --- a/src/SmartSaveImage/nodes/image_saver.py +++ /dev/null @@ -1,332 +0,0 @@ -"""智能图片保存器节点""" - -import os -import json -from datetime import datetime -import folder_paths -import nodes -from ..core import MetadataBuilder, ImageProcessor, PathManager -from ..utils import InputValidator - - -class SmartImageSaver: - """智能图片保存器 - 负责图片保存、格式转换和压缩""" - - CATEGORY = "SmartSave" - OUTPUT_NODE = True - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("images",) - FUNCTION = "save_images" - - # 为翻译插件提供类名标识 - @classmethod - def get_comfy_class(cls): - return "SmartImageSaver" - - def __init__(self): - self.metadata_builder = MetadataBuilder() - self.image_processor = ImageProcessor() - self.path_manager = PathManager() - self.validator = InputValidator() - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "images": ("IMAGE", {"tooltip": "Images to save (usually from SmartFolderManager)"}), - "folder_path": ("STRING", { - "default": "", - "multiline": False, - "tooltip": "Save path (from SmartFolderManager)" - }), - "metadata_json": ("STRING", { - "default": "", - "multiline": True, - "tooltip": "Metadata JSON (from SmartFolderManager)" - }), - "filename_prefix": ("STRING", { - "default": "image", - "multiline": False, - "tooltip": "Filename prefix" - }), - "file_format": (["png", "jpeg", "webp", "bmp", "tiff"], { - "default": "png", - "tooltip": "Image format" - }), - "preview_mode": (["save_and_preview", "preview_only", "save_only"], { - "default": "save_and_preview", - "tooltip": "Save mode" - }), - }, - "optional": { - # 文件名选项 - "add_timestamp": ("BOOLEAN", { - "default": False, - "tooltip": "Add timestamp to filename" - }), - "add_counter": ("BOOLEAN", { - "default": True, - "tooltip": "Add counter (for multiple images)" - }), - "counter_start": ("INT", { - "default": 1, - "min": 0, - "max": 99999, - "tooltip": "Counter start value" - }), - "counter_padding": ("INT", { - "default": 4, - "min": 1, - "max": 10, - "tooltip": "Counter padding (zero-fill)" - }), - - # 图片质量选项 - "jpeg_quality": ("INT", { - "default": 95, - "min": 1, - "max": 100, - "tooltip": "JPEG quality (1-100)" - }), - "webp_quality": ("INT", { - "default": 90, - "min": 1, - "max": 100, - "tooltip": "WebP quality (1-100)" - }), - "webp_lossless": ("BOOLEAN", { - "default": False, - "tooltip": "WebP lossless compression" - }), - "png_compression": ("INT", { - "default": 6, - "min": 0, - "max": 9, - "tooltip": "PNG compression level (0-9)" - }), - - # 元数据选项 - "embed_metadata": ("BOOLEAN", { - "default": True, - "tooltip": "Embed metadata into image" - }), - "embed_workflow": ("BOOLEAN", { - "default": True, - "tooltip": "Embed workflow information" - }), - - # 高级选项 - "overwrite_existing": ("BOOLEAN", { - "default": False, - "tooltip": "Overwrite existing files" - }), - "create_backup": ("BOOLEAN", { - "default": False, - "tooltip": "Create backup for overwritten files" - }), - "optimize_size": ("BOOLEAN", { - "default": False, - "tooltip": "Optimize file size" - }), - }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - }, - } - - def generate_filename(self, prefix, index, add_timestamp, add_counter, - counter_start, counter_padding, file_format): - """生成文件名""" - # 清理前缀 - prefix = self.path_manager.sanitize_filename(prefix) if prefix else "image" - - parts = [prefix] - - # 添加时间戳 - if add_timestamp: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - parts.append(timestamp) - - # 添加计数器 - if add_counter: - counter = str(counter_start + index).zfill(counter_padding) - parts.append(counter) - - # 组合文件名 - filename = "_".join(parts) - - # 添加扩展名 - extensions = { - "png": ".png", - "jpeg": ".jpg", - "webp": ".webp", - "bmp": ".bmp", - "tiff": ".tiff" - } - - filename += extensions.get(file_format, ".png") - return filename - - def parse_metadata_json(self, metadata_json): - """解析元数据JSON""" - if not metadata_json: - return {} - - try: - return json.loads(metadata_json) - except Exception as e: - print(f"[SmartImageSaver] 解析元数据JSON失败: {e}") - return {} - - def save_images(self, images, folder_path, metadata_json, filename_prefix, file_format, preview_mode, - add_timestamp=False, add_counter=True, counter_start=1, counter_padding=4, - jpeg_quality=95, webp_quality=90, webp_lossless=False, png_compression=6, - embed_metadata=True, embed_workflow=False, - overwrite_existing=False, create_backup=False, optimize_size=False, - prompt=None, extra_pnginfo=None): - """保存图片主函数""" - - # 验证输入 - if not self.validator.validate_file_format(file_format): - print(f"[SmartImageSaver] 不支持的文件格式: {file_format}") - file_format = "png" - - if not self.validator.validate_quality_value(jpeg_quality): - jpeg_quality = 95 - - if not self.validator.validate_quality_value(webp_quality): - webp_quality = 90 - - if not self.validator.validate_counter_settings(counter_start, counter_padding): - counter_start, counter_padding = 1, 4 - - # 清理输入 - filename_prefix = self.validator.sanitize_input_string(filename_prefix, 100) - - # 处理预览模式 - if preview_mode == "preview_only": - # 只预览，不保存 - try: - result = nodes.PreviewImage().save_images( - images, - filename_prefix=filename_prefix or "preview" - ) - return {"ui": result.get("ui", {}), "result": (images,)} - except Exception as e: - print(f"[SmartImageSaver] 预览失败: {e}") - return {"ui": {}, "result": (images,)} - - print(f"[SmartImageSaver] 接收到的folder_path: '{folder_path}' (类型: {type(folder_path)})") - folder_path = InputValidator.secure_path_join(folder_path or "", "") - if not os.path.exists(folder_path): - try: - os.makedirs(folder_path, exist_ok=True) - except Exception as e: - folder_path = folder_paths.get_output_directory() - - # 解析元数据 - metadata_dict = self.parse_metadata_json(metadata_json) - workflow_metadata = metadata_dict.get("workflow_metadata", {}) - - # 获取图片尺寸 - if len(images) > 0: - width, height = self.image_processor.get_image_info(images[0]) - workflow_metadata.update({"width": width, "height": height}) - - # 构建元数据文本 - metadata_text = None - if embed_metadata: - metadata_text = self.metadata_builder.build_parameters_text(workflow_metadata) - - # 准备质量设置 - quality_settings = { - "jpeg_quality": jpeg_quality, - "webp_quality": webp_quality, - "webp_lossless": webp_lossless, - "png_compression": png_compression, - "optimize_size": optimize_size, - } - - # 准备工作流数据 - workflow_data = None - if embed_workflow and extra_pnginfo and "workflow" in extra_pnginfo: - workflow_data = extra_pnginfo["workflow"] - - # 保存图片 - saved_images = [] - - for i, image in enumerate(images): - # 生成文件名 - filename = self.generate_filename( - filename_prefix, i, add_timestamp, add_counter, - counter_start, counter_padding, file_format - ) - - # 处理文件名冲突 - if not overwrite_existing: - filename = self.path_manager.generate_unique_filename( - folder_path, - os.path.splitext(filename)[0], - os.path.splitext(filename)[1], - overwrite_existing - ) - - filepath_candidate = os.path.realpath(os.path.join(folder_path, filename)) - base_real = os.path.realpath(folder_path) - try: - if os.path.commonpath([base_real, filepath_candidate]) == base_real: - filepath = filepath_candidate - else: - name_root, name_ext = os.path.splitext(filename) - safe_name = self.path_manager.sanitize_filename(os.path.basename(name_root)) + name_ext - filepath = os.path.realpath(os.path.join(base_real, safe_name)) - except ValueError: - name_root, name_ext = os.path.splitext(filename) - safe_name = self.path_manager.sanitize_filename(os.path.basename(name_root)) + name_ext - filepath = os.path.realpath(os.path.join(folder_paths.get_output_directory(), safe_name)) - - # 创建备份 - if overwrite_existing and create_backup and os.path.exists(filepath): - self.image_processor.create_backup(filepath) - - # 保存图片 - try: - success = self.image_processor.save_image( - image, filepath, file_format, quality_settings, - metadata_text, embed_workflow, workflow_data - ) - - if success: - # 获取相对路径用于UI显示 - try: - rel_folder = os.path.relpath(folder_path, folder_paths.get_output_directory()) - if rel_folder == ".": - rel_folder = "" - except ValueError: - # 如果是绝对路径且不在输出目录下 - rel_folder = os.path.basename(folder_path) - - saved_images.append({ - "filename": filename, - "subfolder": rel_folder, - "type": "output" - }) - print(f"[SmartImageSaver] 保存成功: {filepath}") - else: - print(f"[SmartImageSaver] 保存失败: {filepath}") - - except Exception as e: - print(f"[SmartImageSaver] 保存图片时出错: {e}") - - # 返回结果 - result = {"result": (images,)} - - if preview_mode == "save_and_preview" and saved_images: - # 保存并预览 - result["ui"] = {"images": saved_images} - elif preview_mode == "save_only": - # 仅保存，不显示预览 - 不返回images避免显示预览 - result["ui"] = {} - - return result diff --git a/src/SmartSaveImage/nodes/smart_save.py b/src/SmartSaveImage/nodes/smart_save.py new file mode 100644 index 0000000..129743d --- /dev/null +++ b/src/SmartSaveImage/nodes/smart_save.py @@ -0,0 +1,605 @@ +"""SmartSaveImage —— 规则可视化、目录可预览、模型读取健壮的图片保存节点。 + +设计目标（对应用户诉求）： +1. 所见即所得：保存规则（目录 / 文件名模板）在节点里清晰展示，配合前端实时预览， + 不需要保存完再去猜文件去哪了。 +2. 目录规则强大：支持根目录选择、多级目录模板、丰富的占位符（日期 / 模型 / lora / + 采样器 / 尺寸 / 提示词 / 计数器等），并做安全清洗。 +3. 模型读取健壮：从工作流 prompt 中广度扫描各种加载器、各种输入键、其它插件的自定义 + 节点，尽最大努力拿到 checkpoint / unet / lora / vae 名称，且支持手动覆盖。 + +本文件不依赖任何旧的 core/utils 模块，方便独立维护。 +""" + +from __future__ import annotations + +import json +import os +import re +import tempfile +from datetime import datetime +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +import folder_paths +import nodes as comfy_nodes + +try: + from comfy.cli_args import args +except Exception: # pragma: no cover - 测试环境可能没有 + args = None + + +# --------------------------------------------------------------------------- # +# 常量 +# --------------------------------------------------------------------------- # + +_INVALID_NAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]') +_DATE_TOKEN = re.compile(r"%date(?::([^%]+))?%", re.IGNORECASE) +_WINDOWS_RESERVED = { + "CON", "PRN", "AUX", "NUL", + *(f"COM{i}" for i in range(1, 10)), + *(f"LPT{i}" for i in range(1, 10)), +} + +MODEL_EXTS = (".safetensors", ".ckpt", ".pt", ".pth", ".sft", ".bin", ".gguf", ".onnx") + +# 各类模型输入键（尽量覆盖官方 + 常见第三方节点） +_CKPT_KEYS = ("ckpt_name", "checkpoint", "ckpt", "base_ckpt_name", "model_path") +_UNET_KEYS = ("unet_name", "diffusion_model_name", "model_name") +_LORA_KEYS = ("lora_name", "lora", "lora_name_1", "lora_1", "lora_0") +_VAE_KEYS = ("vae_name", "vae") + +_SAMPLER_SEED_KEYS = ("seed", "noise_seed") + + +# --------------------------------------------------------------------------- # +# 上下文提取（模型 / 种子 / 提示词 / 采样参数） +# --------------------------------------------------------------------------- # + +def _is_link(value: Any) -> bool: + """ComfyUI 中，连线输入会表现为 [node_id, output_index]。""" + return isinstance(value, list) and len(value) == 2 and isinstance(value[0], (str, int)) + + +def _looks_like_model(value: Any) -> bool: + return isinstance(value, str) and value.strip().lower().endswith(MODEL_EXTS) + + +def _iter_prompt_nodes(prompt: Any): + """遍历 prompt（{node_id: {class_type, inputs}}）里的每个节点。""" + if not isinstance(prompt, dict): + return + for node_data in prompt.values(): + if isinstance(node_data, dict) and "class_type" in node_data: + inputs = node_data.get("inputs") + yield str(node_data.get("class_type", "")), inputs if isinstance(inputs, dict) else {} + + +def extract_context(prompt: Any, width: int = 0, height: int = 0) -> dict[str, Any]: + """从工作流 prompt 中健壮地提取用于命名的上下文信息。 + + 返回值一定包含所有键，缺失时给出合理的占位值，绝不抛异常。 + """ + ctx: dict[str, Any] = { + "model": "", + "model_full": "", + "unet": "", + "lora": "", + "loras": [], + "vae": "", + "seed": "", + "steps": "", + "cfg": "", + "sampler": "", + "scheduler": "", + "positive": "", + "negative": "", + "width": str(int(width)) if width else "", + "height": str(int(height)) if height else "", + } + + checkpoint = "" + unet = "" + vae = "" + loras: list[str] = [] + text_candidates: list[str] = [] + + for class_type, inputs in _iter_prompt_nodes(prompt): + cls_lower = class_type.lower() + + # ---- 模型：checkpoint / unet ---- + for key, value in inputs.items(): + if _is_link(value): + continue + key_l = key.lower() + if not checkpoint and (key_l in _CKPT_KEYS or (_looks_like_model(value) and "ckpt" in key_l)): + if isinstance(value, str) and value.strip(): + checkpoint = value.strip() + if not unet and key_l in _UNET_KEYS and isinstance(value, str) and value.strip(): + unet = value.strip() + if not vae and key_l in _VAE_KEYS and isinstance(value, str) and value.strip() and value != "taesd": + vae = value.strip() + # lora 可能有多个键（lora_name / lora_1 / lora_02 ...） + if isinstance(value, str) and value.strip() and ("lora" in key_l) and _looks_like_model(value): + loras.append(value.strip()) + + # 兜底：某些加载器把模型放在不常见的键里，但类名里带 Checkpoint/Loader + if not checkpoint and ("checkpoint" in cls_lower or "ckpt" in cls_lower): + for value in inputs.values(): + if not _is_link(value) and _looks_like_model(value): + checkpoint = value.strip() + break + + # ---- 采样参数 ---- + if "sampler" in cls_lower or "ksampler" in cls_lower: + for k in _SAMPLER_SEED_KEYS: + v = inputs.get(k) + if not ctx["seed"] and isinstance(v, (int, float, str)) and not _is_link(v): + ctx["seed"] = str(int(v)) if isinstance(v, float) else str(v) + for field in ("steps", "cfg", "sampler_name", "scheduler"): + v = inputs.get(field) + if _is_link(v): + continue + target = {"steps": "steps", "cfg": "cfg", + "sampler_name": "sampler", "scheduler": "scheduler"}[field] + if not ctx[target] and isinstance(v, (int, float, str)) and str(v) != "": + ctx[target] = str(v) + + # ---- 提示词文本 ---- + if "textencode" in cls_lower or "cliptextencode" in cls_lower or "text" in cls_lower: + for tkey in ("text", "text_g", "text_l", "prompt", "positive"): + v = inputs.get(tkey) + if isinstance(v, str) and v.strip() and not _is_link(v): + text_candidates.append(v.strip()) + + # ---- 尺寸兜底（预览时没有真实图片，从 EmptyLatentImage 等读取）---- + if not ctx["width"]: + w = inputs.get("width") + if isinstance(w, (int, float)) and not _is_link(w): + ctx["width"] = str(int(w)) + if not ctx["height"]: + h = inputs.get("height") + if isinstance(h, (int, float)) and not _is_link(h): + ctx["height"] = str(int(h)) + + # 模型主名：优先 checkpoint，其次 unet + model_full = checkpoint or unet + ctx["model_full"] = model_full + ctx["model"] = Path(model_full).stem if model_full else "" + ctx["unet"] = Path(unet).stem if unet else "" + ctx["vae"] = Path(vae).stem if vae else "" + # 去重保序 + seen = set() + uniq_loras = [] + for l in loras: + stem = Path(l).stem + if stem not in seen: + seen.add(stem) + uniq_loras.append(stem) + ctx["loras"] = uniq_loras + ctx["lora"] = uniq_loras[0] if uniq_loras else "" + + if text_candidates: + ctx["positive"] = text_candidates[0] + if len(text_candidates) > 1: + ctx["negative"] = text_candidates[1] + + return ctx + + +# --------------------------------------------------------------------------- # +# 模板引擎 +# --------------------------------------------------------------------------- # + +def _strftime_format(value: str) -> str: + result = value + for token, repl in ( + ("yyyy", "%Y"), ("yy", "%y"), ("MM", "%m"), ("dd", "%d"), + ("HH", "%H"), ("hh", "%H"), ("mm", "%M"), ("ss", "%S"), + ): + result = result.replace(token, repl) + return result + + +def sanitize_segment(value: str, fallback: str = "", max_length: int = 80) -> str: + value = _INVALID_NAME_CHARS.sub("_", str(value)).strip().strip(".") + value = re.sub(r"\s+", " ", value) + if value.upper() in _WINDOWS_RESERVED: + value = f"_{value}" + value = value[:max_length].rstrip() + return value or fallback + + +def expand_template(template: str, ctx: dict[str, Any], batch_index: int = 0, + prompt_max_len: int = 60) -> str: + """把模板字符串中的占位符替换为真实值（不做路径分段清洗）。""" + now = datetime.now() + + def replace_date(match: re.Match[str]) -> str: + fmt = match.group(1) or "yyyy-MM-dd" + try: + return now.strftime(_strftime_format(fmt)) + except ValueError: + return now.strftime("%Y-%m-%d") + + value = _DATE_TOKEN.sub(replace_date, template or "") + + positive = re.sub(r"\s+", " ", ctx.get("positive", "")).strip()[:prompt_max_len] + + # Only literal separators typed in the template may create directories. + # Values coming from models/prompts often contain slashes and must stay in one segment. + model = sanitize_segment(ctx.get("model") or "unknown_model", "unknown_model", 100) + model_full = sanitize_segment(ctx.get("model_full") or "unknown_model", "unknown_model", 140) + unet = sanitize_segment(ctx.get("unet") or "", "", 100) + lora = sanitize_segment(ctx.get("lora") or "no_lora", "no_lora", 100) + vae = sanitize_segment(ctx.get("vae") or "", "", 100) + prompt_text = sanitize_segment(positive or "untitled", "untitled", prompt_max_len) + + tokens = { + "%year%": now.strftime("%Y"), + "%month%": now.strftime("%m"), + "%day%": now.strftime("%d"), + "%hour%": now.strftime("%H"), + "%minute%": now.strftime("%M"), + "%second%": now.strftime("%S"), + "%model%": model, + "%model_full%": model_full, + "%unet%": unet, + "%lora%": lora, + "%vae%": vae, + "%seed%": str(ctx.get("seed") or "0"), + "%steps%": str(ctx.get("steps") or ""), + "%cfg%": str(ctx.get("cfg") or ""), + "%sampler%": ctx.get("sampler") or "", + "%scheduler%": ctx.get("scheduler") or "", + "%width%": str(ctx.get("width") or "0"), + "%height%": str(ctx.get("height") or "0"), + "%prompt%": prompt_text, + "%batch%": f"{batch_index:02d}", + } + for token, repl in tokens.items(): + value = value.replace(token, str(repl)) + return value + + +def build_subfolder(template: str, ctx: dict[str, Any], prompt_max_len: int = 60) -> str: + """把目录模板展开成安全的相对子目录（保留多级）。""" + expanded = expand_template(template, ctx, 0, prompt_max_len) + parts = re.split(r"[\\/]+", expanded) + safe = [sanitize_segment(p) for p in parts if p.strip() not in {"", ".", ".."}] + safe = [p for p in safe if p] + return os.path.join(*safe) if safe else "" + + +def build_filename_base(template: str, ctx: dict[str, Any], batch_index: int, + batch_size: int, counter_digits: int) -> str: + """Build one safe filename stem using the same rules for preview and saving.""" + expanded = expand_template(template or "image", ctx, batch_index) + base = sanitize_segment(expanded, fallback="image", max_length=160) + if batch_size > 1 and "%batch%" not in (template or ""): + width = max(int(counter_digits), 1) + base = f"{base}_{batch_index:0{width}d}" + return base + + +# --------------------------------------------------------------------------- # +# 根目录解析 +# --------------------------------------------------------------------------- # + +def _known_roots() -> dict[str, str]: + roots = { + "output": folder_paths.get_output_directory(), + "input": folder_paths.get_input_directory(), + "temp": folder_paths.get_temp_directory(), + } + return {k: os.path.realpath(v) for k, v in roots.items()} + + +def resolve_root(root_mode: str, custom_root: str) -> str: + roots = _known_roots() + if root_mode in roots: + return roots[root_mode] + # custom + output_root = roots["output"] + if not custom_root or not custom_root.strip(): + return output_root + requested = os.path.expandvars(os.path.expanduser(custom_root.strip())) + if not os.path.isabs(requested): + requested = os.path.join(output_root, requested) + return os.path.realpath(requested) + + +def resolve_target(root_mode: str, custom_root: str, folder_template: str, + ctx: dict[str, Any], prompt_max_len: int = 60) -> tuple[str, str]: + """返回 (根目录, 完整目标目录)。""" + root = resolve_root(root_mode, custom_root) + subfolder = build_subfolder(folder_template, ctx, prompt_max_len) + target = os.path.realpath(os.path.join(root, subfolder)) + return root, target + + +# --------------------------------------------------------------------------- # +# 节点主体 +# --------------------------------------------------------------------------- # + +class SmartSaveImage: + """规则可视化的智能图片保存节点。""" + + CATEGORY = "image/save" + DESCRIPTION = "所见即所得的图片保存：目录/文件名规则清晰可见，支持实时目录预览与健壮的模型读取。" + SEARCH_ALIASES = ["save image", "smart save", "保存图片", "智能保存"] + OUTPUT_NODE = True + RETURN_TYPES = () + FUNCTION = "save_images" + + ROOT_MODES = ["output", "custom", "input", "temp"] + FORMATS = ["png", "jpeg", "webp"] + COLLISION_INCREMENT = "increment" + COLLISION_OVERWRITE = "overwrite" + MODE_SAVE_PREVIEW = "save_and_preview" + MODE_SAVE_ONLY = "save_only" + MODE_PREVIEW_ONLY = "preview_only" + + @classmethod + def INPUT_TYPES(cls): + try: + ckpts = folder_paths.get_filename_list("checkpoints") + except Exception: + ckpts = [] + model_choices = ["auto"] + ckpts + + return { + "required": { + "images": ("IMAGE", {"tooltip": "需要保存的图片或图片批次。"}), + "root_mode": (cls.ROOT_MODES, { + "default": "output", + "tooltip": "根目录：output=ComfyUI 输出目录；custom=自定义绝对/相对路径；input/temp 为内置目录。", + }), + "custom_root": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "仅当 root_mode=custom 生效；可用绝对路径，留空回退到 output。", + }), + "folder_template": ("STRING", { + "default": "%date:yyyy-MM-dd%/%model%", + "multiline": False, + "tooltip": "子目录模板，支持多级与占位符，如 %date%/%model%/%sampler%。", + }), + "filename_template": ("STRING", { + "default": "%model%_%seed%", + "multiline": False, + "tooltip": "文件名模板（不含扩展名）。批量时自动追加序号，或用 %batch% 自定义位置。", + }), + "file_format": (cls.FORMATS, {"default": "png"}), + "quality": ("INT", { + "default": 95, "min": 1, "max": 100, "step": 1, + "tooltip": "旧工作流兼容字段，不在新版界面中显示。", + }), + "collision_mode": ([cls.COLLISION_INCREMENT, cls.COLLISION_OVERWRITE], { + "default": cls.COLLISION_INCREMENT, + "tooltip": "increment=自动编号避免覆盖；overwrite=覆盖同名文件。", + }), + "save_mode": ([cls.MODE_SAVE_PREVIEW, cls.MODE_SAVE_ONLY, cls.MODE_PREVIEW_ONLY], { + "default": cls.MODE_SAVE_PREVIEW, + }), + }, + "optional": { + "manual_model": (model_choices, { + "default": "auto", + "tooltip": "手动指定模型名以覆盖自动读取（auto=自动从工作流读取）。", + }), + "embed_workflow": ("BOOLEAN", { + "default": True, + "tooltip": "PNG 写入 ComfyUI 工作流元数据；JPEG/WebP 写入 EXIF。", + }), + "counter_digits": ("INT", { + "default": 3, "min": 0, "max": 8, + "tooltip": "批量/防冲突序号位数，0 表示尽量不加序号。", + }), + "png_compression": ("INT", { + "default": 4, "min": 0, "max": 9, "step": 1, + "tooltip": "PNG 无损压缩等级；4 与 ComfyUI 自带保存节点一致。", + }), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + # --------------------------- 元数据 --------------------------- # + + @staticmethod + def _safe_json(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str) + + @staticmethod + def _metadata_disabled() -> bool: + return args is not None and getattr(args, "disable_metadata", False) + + @classmethod + def _png_metadata(cls, prompt, extra_pnginfo, ctx, enabled: bool) -> PngInfo | None: + if not enabled or cls._metadata_disabled(): + return None + meta = PngInfo() + # A1111 风格 parameters，方便第三方查看器识别 + params = cls._build_parameters_text(ctx) + if params: + meta.add_text("parameters", params) + if prompt is not None: + meta.add_text("prompt", cls._safe_json(prompt)) + if isinstance(extra_pnginfo, dict): + for key, value in extra_pnginfo.items(): + meta.add_text(str(key), cls._safe_json(value)) + return meta + + @classmethod + def _exif_metadata(cls, image: Image.Image, prompt, extra_pnginfo, enabled: bool): + if not enabled or cls._metadata_disabled(): + return None + payload = {"prompt": prompt, "extra_pnginfo": extra_pnginfo} + exif = image.getexif() + exif[0x9286] = cls._safe_json(payload) # UserComment + return exif + + @staticmethod + def _build_parameters_text(ctx: dict[str, Any]) -> str: + parts: list[str] = [] + if ctx.get("positive"): + parts.append(ctx["positive"]) + if ctx.get("negative"): + parts.append(f"Negative prompt: {ctx['negative']}") + extra = [] + if ctx.get("steps"): + extra.append(f"Steps: {ctx['steps']}") + if ctx.get("sampler"): + sch = ctx.get("scheduler") + extra.append(f"Sampler: {ctx['sampler']}" + (f" {sch}" if sch and sch != "normal" else "")) + if ctx.get("cfg"): + extra.append(f"CFG scale: {ctx['cfg']}") + if ctx.get("seed"): + extra.append(f"Seed: {ctx['seed']}") + if ctx.get("width") and ctx.get("height"): + extra.append(f"Size: {ctx['width']}x{ctx['height']}") + if ctx.get("model"): + extra.append(f"Model: {ctx['model']}") + if extra: + parts.append(", ".join(extra)) + return "\n".join(parts) + + # --------------------------- 图片写盘 --------------------------- # + + @staticmethod + def _extension(file_format: str) -> str: + return {"png": ".png", "jpeg": ".jpg", "webp": ".webp"}.get(file_format, ".png") + + @staticmethod + def _tensor_to_image(image) -> Image.Image: + array = image.detach().cpu().numpy() if hasattr(image, "detach") else np.asarray(image) + array = np.clip(array * 255.0, 0, 255).astype(np.uint8) + return Image.fromarray(array) + + @classmethod + def _save_one(cls, image, path, file_format, quality, png_compression, + prompt, extra_pnginfo, ctx, embed): + pil = cls._tensor_to_image(image) + pil_format = {"png": "PNG", "jpeg": "JPEG", "webp": "WEBP"}.get(file_format, "PNG") + kwargs: dict[str, Any] = {} + if file_format == "png": + kwargs["pnginfo"] = cls._png_metadata(prompt, extra_pnginfo, ctx, embed) + kwargs["compress_level"] = max(0, min(int(png_compression), 9)) + else: + if pil.mode not in {"RGB", "L"}: + pil = pil.convert("RGB") + kwargs["exif"] = cls._exif_metadata(pil, prompt, extra_pnginfo, embed) + if file_format == "webp": + kwargs["lossless"] = True + else: + kwargs["quality"] = max(1, min(int(quality), 100)) + kwargs["subsampling"] = 0 + + os.makedirs(os.path.dirname(path), exist_ok=True) + temp_path = "" + try: + with tempfile.NamedTemporaryFile(prefix=".smartsave-", suffix=cls._extension(file_format), + dir=os.path.dirname(path), delete=False) as handle: + temp_path = handle.name + pil.save(temp_path, format=pil_format, + **{k: v for k, v in kwargs.items() if v is not None}) + os.replace(temp_path, path) + temp_path = "" + finally: + if temp_path and os.path.exists(temp_path): + os.remove(temp_path) + + @staticmethod + def _unique_path(path: str, overwrite: bool, digits: int) -> str: + if overwrite or not os.path.exists(path): + return path + root, ext = os.path.splitext(path) + width = max(digits, 1) + counter = 1 + while counter <= 9_999_999: + candidate = f"{root}_{counter:0{width}d}{ext}" + if not os.path.exists(candidate): + return candidate + counter += 1 + raise RuntimeError(f"无法为文件生成可用编号: {path}") + + @staticmethod + def _is_inside(path: str, root: str) -> bool: + try: + return os.path.commonpath([os.path.realpath(path), os.path.realpath(root)]) == os.path.realpath(root) + except ValueError: + return False + + @classmethod + def _ui_image(cls, path: str) -> dict[str, str] | None: + roots = ( + ("output", folder_paths.get_output_directory()), + ("input", folder_paths.get_input_directory()), + ("temp", folder_paths.get_temp_directory()), + ) + for image_type, root in roots: + if not cls._is_inside(path, root): + continue + subfolder = os.path.relpath(os.path.dirname(path), root) + return { + "filename": os.path.basename(path), + "subfolder": "" if subfolder == "." else subfolder.replace("\\", "/"), + "type": image_type, + } + return None + + # --------------------------- 主函数 --------------------------- # + + def save_images(self, images, root_mode, custom_root, folder_template, filename_template, + file_format, quality, collision_mode, save_mode, + manual_model="auto", embed_workflow=True, counter_digits=3, png_compression=4, + prompt=None, extra_pnginfo=None): + + if save_mode == self.MODE_PREVIEW_ONLY: + return comfy_nodes.PreviewImage().save_images( + images, filename_prefix="SmartSavePreview", + prompt=prompt, extra_pnginfo=extra_pnginfo) + + height = int(images[0].shape[0]) if len(images) else 0 + width = int(images[0].shape[1]) if len(images) else 0 + ctx = extract_context(prompt, width, height) + if manual_model and manual_model != "auto": + ctx["model_full"] = manual_model + ctx["model"] = Path(manual_model).stem + + root, target_folder = resolve_target(root_mode, custom_root, folder_template, ctx) + os.makedirs(target_folder, exist_ok=True) + + overwrite = collision_mode == self.COLLISION_OVERWRITE + saved_ui: list[dict[str, str]] = [] + + for batch_index, image in enumerate(images): + base = build_filename_base( + filename_template, ctx, batch_index, len(images), counter_digits) + filename = base + self._extension(file_format) + path = self._unique_path(os.path.join(target_folder, filename), overwrite, counter_digits) + + self._save_one( + image, path, file_format, quality, png_compression, + prompt, extra_pnginfo, ctx, embed_workflow) + + ui_image = self._ui_image(path) + if ui_image: + saved_ui.append(ui_image) + + if save_mode == self.MODE_SAVE_ONLY: + return {"ui": {}} + if saved_ui: + return {"ui": {"images": saved_ui}} + # 保存到 output 之外时无法通过 /view 直接预览，用临时预览兜底 + return comfy_nodes.PreviewImage().save_images( + images, filename_prefix="SmartSavePreview", + prompt=prompt, extra_pnginfo=extra_pnginfo) diff --git a/src/SmartSaveImage/server_routes.py b/src/SmartSaveImage/server_routes.py new file mode 100644 index 0000000..c8e41c3 --- /dev/null +++ b/src/SmartSaveImage/server_routes.py @@ -0,0 +1,125 @@ +"""为 SmartSaveImage 提供实时目录预览的后端路由。 + +前端把当前工作流 prompt 和模板发过来，这里用与保存时完全相同的解析逻辑算出： +- 解析到的上下文（model / seed / sampler ...），让用户一眼看清占位符会变成什么； +- 最终目标目录、示例文件名； +- 目录是否已存在、已有多少张图片。 + +这样用户在保存前就能确认规则，而不是保存完再去猜。 +""" + +from __future__ import annotations + +import os + +from .nodes.smart_save import ( + SmartSaveImage, + build_filename_base, + build_subfolder, + extract_context, + resolve_root, +) + +_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".gif") + + +def _count_images(folder: str) -> int: + try: + return sum( + 1 for name in os.listdir(folder) + if name.lower().endswith(_IMAGE_EXTS) and os.path.isfile(os.path.join(folder, name)) + ) + except Exception: + return 0 + + +def compute_preview(data: dict) -> dict: + prompt = data.get("prompt") or {} + root_mode = data.get("root_mode", "output") + custom_root = data.get("custom_root", "") + folder_template = data.get("folder_template", "") + filename_template = data.get("filename_template", "image") + file_format = data.get("file_format", "png") + manual_model = data.get("manual_model", "auto") + counter_digits = min(max(int(data.get("counter_digits", 3) or 0), 0), 8) + batch_size = min(max(int(data.get("batch_size", 1) or 1), 1), 1000) + collision_mode = data.get("collision_mode", SmartSaveImage.COLLISION_INCREMENT) + + ctx = extract_context(prompt) + if manual_model and manual_model != "auto": + from pathlib import Path + ctx["model_full"] = manual_model + ctx["model"] = Path(manual_model).stem + + root = resolve_root(root_mode, custom_root) + subfolder = build_subfolder(folder_template, ctx) + target = os.path.realpath(os.path.join(root, subfolder)) + + ext = SmartSaveImage._extension(file_format) + example_names: list[str] = [] + preview_count = min(max(batch_size, 1), 3) + for i in range(preview_count): + base = build_filename_base(filename_template, ctx, i, batch_size, counter_digits) + desired_path = os.path.join(target, base + ext) + preview_path = SmartSaveImage._unique_path( + desired_path, + collision_mode == SmartSaveImage.COLLISION_OVERWRITE, + counter_digits, + ) + example_names.append(os.path.basename(preview_path)) + + exists = os.path.isdir(target) + return { + "ok": True, + "root": root, + "subfolder": subfolder, + "target": target, + "exists": exists, + "existing_count": _count_images(target) if exists else 0, + "example_filenames": example_names, + "context": { + "model": ctx.get("model") or "", + "model_full": ctx.get("model_full") or "", + "lora": ctx.get("lora") or "", + "loras": ctx.get("loras", []), + "vae": ctx.get("vae") or "", + "seed": ctx.get("seed") or "", + "steps": ctx.get("steps") or "", + "cfg": ctx.get("cfg") or "", + "sampler": ctx.get("sampler") or "", + "scheduler": ctx.get("scheduler") or "", + "width": ctx.get("width") or "", + "height": ctx.get("height") or "", + "positive": (ctx.get("positive") or "")[:120], + }, + } + + +def register_routes() -> None: + try: + from server import PromptServer + from aiohttp import web + except Exception: + return + + instance = getattr(PromptServer, "instance", None) + if instance is None or not hasattr(instance, "routes"): + return + + # 避免重复注册 + if getattr(instance, "_smartsave_routes_registered", False): + return + + @instance.routes.post("/smartsave/preview") + async def smartsave_preview(request): # noqa: ANN001 + try: + data = await request.json() + except Exception: + data = {} + try: + result = compute_preview(data or {}) + except Exception as exc: # 预览失败不影响使用 + result = {"ok": False, "error": str(exc)} + return web.json_response(result) + + instance._smartsave_routes_registered = True diff --git a/src/SmartSaveImage/utils/__init__.py b/src/SmartSaveImage/utils/__init__.py deleted file mode 100644 index c5f5ea8..0000000 --- a/src/SmartSaveImage/utils/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""SmartSaveImage 工具模块""" - -from .validators import InputValidator - -__all__ = [ - "InputValidator", -] diff --git a/src/SmartSaveImage/utils/validators.py b/src/SmartSaveImage/utils/validators.py deleted file mode 100644 index a9e0a1a..0000000 --- a/src/SmartSaveImage/utils/validators.py +++ /dev/null @@ -1,147 +0,0 @@ -"""输入验证工具""" - -import os -import re -from typing import Any, List, Optional, Union - - -class InputValidator: - """输入验证器""" - - @staticmethod - def validate_folder_path(path: str, allowed_base_paths: list = None) -> bool: - """验证文件夹路径是否有效且安全""" - if not path: - return False - - # 检查路径长度 - if len(path) > 260: # Windows路径长度限制 - return False - - # 检查非法字符 - illegal_chars = r'[<>"|?*]' - if re.search(illegal_chars, path): - return False - - # 检查路径遍历攻击 - if '..' in path or path.startswith('/') or '~' in path: - return False - - return True - - @staticmethod - def secure_path_join(base_path: str, user_path: str = "") -> str: - """安全地连接路径，防止路径遍历攻击 - 使用GitHub建议的方法""" - import folder_paths - - # 获取ComfyUI的输出目录作为安全基础路径 - safe_base = folder_paths.get_output_directory() - - try: - # 使用realpath解析安全基础路径 - resolved_safe_base = os.path.realpath(safe_base) - - # 处理base_path - if os.path.isabs(base_path): - # 绝对路径：直接解析 - candidate_path = os.path.realpath(base_path) - else: - # 相对路径：相对于安全基础路径 - candidate_path = os.path.realpath(os.path.join(safe_base, base_path)) - - # 如果有用户路径，继续连接 - if user_path: - candidate_path = os.path.realpath(os.path.join(candidate_path, user_path)) - - # 使用commonpath确保路径在安全范围内 - try: - common = os.path.commonpath([resolved_safe_base, candidate_path]) - # 检查公共路径是否就是安全基础路径 - if os.path.realpath(common) == resolved_safe_base: - return candidate_path - else: - # 路径逃逸，返回安全基础路径 - return resolved_safe_base - except ValueError: - # commonpath失败（通常是不同驱动器），返回安全基础路径 - return resolved_safe_base - - except (OSError, ValueError) as e: - # 任何路径操作失败，返回安全基础路径 - return safe_base - - @staticmethod - def validate_filename(filename: str) -> bool: - """验证文件名是否有效""" - if not filename: - return False - - # 检查长度 - if len(filename) > 255: - return False - - # 检查非法字符 - illegal_chars = r'[<>:"/\\|?*]' - if re.search(illegal_chars, filename): - return False - - # 检查保留名称（Windows） - reserved_names = [ - 'CON', 'PRN', 'AUX', 'NUL', - 'COM1', 'COM2', 'COM3', 'COM4', 'COM5', 'COM6', 'COM7', 'COM8', 'COM9', - 'LPT1', 'LPT2', 'LPT3', 'LPT4', 'LPT5', 'LPT6', 'LPT7', 'LPT8', 'LPT9' - ] - - name_without_ext = os.path.splitext(filename)[0].upper() - if name_without_ext in reserved_names: - return False - - return True - - @staticmethod - def validate_date_format(date_format: str) -> bool: - """验证日期格式字符串""" - if not date_format: - return False - - # 检查是否包含有效的日期格式标记 - valid_tokens = ['yyyy', 'yy', 'MM', 'dd', 'hh', 'mm', 'ss'] - has_valid_token = any(token in date_format for token in valid_tokens) - - return has_valid_token - - @staticmethod - def validate_quality_value(quality: int, min_val: int = 1, max_val: int = 100) -> bool: - """验证质量值范围""" - return isinstance(quality, int) and min_val <= quality <= max_val - - @staticmethod - def validate_file_format(file_format: str) -> bool: - """验证文件格式""" - supported_formats = ['png', 'jpeg', 'webp', 'bmp', 'tiff'] - return file_format.lower() in supported_formats - - @staticmethod - def sanitize_input_string(input_str: str, max_length: int = 1000) -> str: - """清理输入字符串""" - if not isinstance(input_str, str): - input_str = str(input_str) - - # 移除控制字符 - input_str = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', input_str) - - # 限制长度 - if len(input_str) > max_length: - input_str = input_str[:max_length] - - return input_str.strip() - - @staticmethod - def validate_counter_settings(counter_start: int, counter_padding: int) -> bool: - """验证计数器设置""" - return ( - isinstance(counter_start, int) and - isinstance(counter_padding, int) and - 0 <= counter_start <= 99999 and - 1 <= counter_padding <= 10 - ) diff --git a/tests/pytest.ini b/tests/pytest.ini index 95c76f1..375278d 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -1,4 +1,4 @@ [pytest] -testpaths = . # Run tests in the current directory -python_files = test_*.py # Run tests in files that start with "test_" -norecursedirs = .. # Don't run tests in the parent directory +testpaths = . +python_files = test_*.py +norecursedirs = __pycache__ diff --git a/tests/test_SmartSaveImage.py b/tests/test_SmartSaveImage.py index 80d6660..bf70d22 100644 --- a/tests/test_SmartSaveImage.py +++ b/tests/test_SmartSaveImage.py @@ -1,21 +1,191 @@ -#!/usr/bin/env python +"""SmartSaveImage 节点回归测试。""" -"""Tests for `SmartSaveImage` package.""" +from pathlib import Path import pytest -from src.SmartSaveImage.nodes import Example +import torch +from PIL import Image -@pytest.fixture -def example_node(): - """Fixture to create an Example node instance.""" - return Example() +pytest.importorskip("folder_paths") -def test_example_node_initialization(example_node): - """Test that the node can be instantiated.""" - assert isinstance(example_node, Example) +from src.SmartSaveImage.nodes import NODE_CLASS_MAPPINGS, SmartSaveImage +from src.SmartSaveImage.nodes.smart_save import ( + build_subfolder, + expand_template, + extract_context, +) +from src.SmartSaveImage.server_routes import compute_preview -def test_return_types(): - """Test the node's metadata.""" - assert Example.RETURN_TYPES == ("IMAGE",) - assert Example.FUNCTION == "test" - assert Example.CATEGORY == "Example" + +def test_node_is_registered(): + assert NODE_CLASS_MAPPINGS["SmartSaveImage"] is SmartSaveImage + assert SmartSaveImage.OUTPUT_NODE is True + assert SmartSaveImage.RETURN_TYPES == () + + +def test_template_expansion_and_sanitizing(): + context = { + "model": "model:name", + "seed": "42", + "positive": "portrait / studio", + "width": "1024", + "height": "768", + } + + folder = build_subfolder("%model%/%seed%/%prompt%", context) + + assert Path(folder).parts == ("model_name", "42", "portrait _ studio") + + +def test_extracts_model_name_from_unet_loader(): + prompt = { + "1": { + "class_type": "UNETLoader", + "inputs": {"unet_name": "Krea/Krea2_fp8.safetensors"}, + } + } + + context = extract_context(prompt) + + assert context["model"] == "Krea2_fp8" + + +def test_extract_context_ignores_links_and_reads_sampler(): + prompt = { + "1": {"class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "sdxl/dreamShaperXL.safetensors"}}, + "2": {"class_type": "LoraLoader", + "inputs": {"lora_name": "add_detail.safetensors", "model": ["1", 0]}}, + "3": {"class_type": "KSampler", + "inputs": {"seed": 999, "steps": 30, "cfg": 6.0, + "sampler_name": "euler", "scheduler": "normal", "model": ["2", 0]}}, + } + + ctx = extract_context(prompt) + + assert ctx["model"] == "dreamShaperXL" + assert ctx["lora"] == "add_detail" + assert ctx["seed"] == "999" + assert ctx["sampler"] == "euler" + + +def test_collision_mode_adds_counter(tmp_path): + original = tmp_path / "image.png" + original.write_bytes(b"existing") + + result = SmartSaveImage._unique_path(str(original), overwrite=False, digits=3) + + assert Path(result).name == "image_001.png" + assert original.read_bytes() == b"existing" + + +def test_preview_uses_next_available_filename(tmp_path): + (tmp_path / "image.png").write_bytes(b"existing") + + preview = compute_preview({ + "root_mode": "custom", + "custom_root": str(tmp_path), + "folder_template": "", + "filename_template": "image", + "file_format": "png", + "collision_mode": "increment", + "counter_digits": 3, + }) + + assert preview["example_filenames"] == ["image_001.png"] + + +def test_batch_token_in_filename(): + ctx = extract_context({}) + assert expand_template("img_%batch%", ctx, 5) == "img_05" + + +@pytest.mark.parametrize("file_format,pil_format", [ + ("png", "PNG"), + ("jpeg", "JPEG"), + ("webp", "WEBP"), +]) +def test_saves_real_image_batches(tmp_path, file_format, pil_format): + images = torch.rand((2, 16, 20, 3)) + node = SmartSaveImage() + + node.save_images( + images=images, + root_mode="custom", + custom_root=str(tmp_path), + folder_template="case_%model%", + filename_template="sample_%batch%", + file_format=file_format, + quality=91, + collision_mode=SmartSaveImage.COLLISION_INCREMENT, + save_mode=SmartSaveImage.MODE_SAVE_ONLY, + manual_model="auto", + embed_workflow=True, + counter_digits=3, + prompt={ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "models/demo.safetensors"}, + } + }, + extra_pnginfo={"workflow": {"nodes": []}}, + ) + + ext = SmartSaveImage._extension(file_format).lstrip(".") + saved = sorted(tmp_path.rglob(f"*.{ext}")) + assert [path.stem for path in saved] == ["sample_00", "sample_01"] + assert {Image.open(path).format for path in saved} == {pil_format} + + +def test_overwrite_batch_with_zero_digits_keeps_every_image(tmp_path): + images = torch.rand((2, 8, 8, 3)) + + SmartSaveImage().save_images( + images=images, + root_mode="custom", + custom_root=str(tmp_path), + folder_template="", + filename_template="image", + file_format="png", + quality=95, + collision_mode=SmartSaveImage.COLLISION_OVERWRITE, + save_mode=SmartSaveImage.MODE_SAVE_ONLY, + manual_model="auto", + embed_workflow=False, + counter_digits=0, + ) + + assert sorted(path.name for path in tmp_path.glob("*.png")) == ["image_0.png", "image_1.png"] + + +def test_png_compression_defaults_to_comfyui_save_image_level(): + config = SmartSaveImage.INPUT_TYPES()["optional"]["png_compression"][1] + + assert config["default"] == 4 + assert config["min"] == 0 + assert config["max"] == 9 + + +def test_png_compression_changes_size_without_changing_pixels(tmp_path): + image = torch.zeros((1, 64, 64, 3)) + node = SmartSaveImage() + common = { + "images": image, + "root_mode": "custom", + "custom_root": str(tmp_path), + "folder_template": "", + "file_format": "png", + "quality": 95, + "collision_mode": SmartSaveImage.COLLISION_OVERWRITE, + "save_mode": SmartSaveImage.MODE_SAVE_ONLY, + "embed_workflow": False, + } + + node.save_images(filename_template="level_0", png_compression=0, **common) + node.save_images(filename_template="level_4", png_compression=4, **common) + + level_0 = tmp_path / "level_0.png" + level_4 = tmp_path / "level_4.png" + with Image.open(level_0) as first, Image.open(level_4) as second: + assert first.tobytes() == second.tobytes() + assert level_4.stat().st_size < level_0.stat().st_size diff --git a/web/smart_save.css b/web/smart_save.css new file mode 100644 index 0000000..6ff95cf --- /dev/null +++ b/web/smart_save.css @@ -0,0 +1,270 @@ +.ssi-panel { + display: flex; + flex-direction: column; + gap: 9px; + padding: 6px 8px 10px; + box-sizing: border-box; + width: 100%; + min-width: 0; + font-size: 12px; + color: var(--input-text, #e6e6e6); +} + +.ssi-row { + display: flex; + align-items: center; + gap: 8px; +} + +.ssi-field { + display: flex; + flex-direction: column; + gap: 3px; + min-width: 0; +} + +.ssi-label { + font-weight: 600; + color: var(--descrip-text, #b8c0cc); + min-width: 60px; +} + +.ssi-input, +.ssi-select { + flex: 1; + width: 100%; + min-width: 0; + height: 30px; + box-sizing: border-box; + background: var(--comfy-input-bg, #222); + color: var(--input-text, #e6e6e6); + border: 1px solid var(--border-color, #444); + border-radius: 5px; + padding: 5px 8px; + outline: none; +} + +.ssi-input:focus, +.ssi-select:focus { + border-color: #5b9dff; +} + +.ssi-mono { + font-family: "Consolas", "Menlo", monospace; + letter-spacing: 0.2px; +} + +/* 根目录分段按钮 */ +.ssi-location { + align-items: stretch; +} + +.ssi-seg { + display: grid; + grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 3px; + width: 100%; +} + +.ssi-seg-btn { + min-width: 0; + height: 30px; + padding: 4px 3px; + border: 1px solid var(--border-color, #444); + border-radius: 5px; + background: var(--comfy-input-bg, #222); + color: var(--descrip-text, #b8c0cc); + cursor: pointer; + transition: all 0.12s ease; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.ssi-seg-btn:hover { + border-color: #5b9dff; +} + +.ssi-seg-active { + background: #3a6df0; + color: #fff; + border-color: #3a6df0; +} + +/* 占位符调色板 */ +.ssi-palette { + padding: 7px 8px; + background: rgba(255, 255, 255, 0.03); + border: 1px solid var(--border-color, #444); + border-radius: 6px; +} + +.ssi-palette-summary { + cursor: pointer; + color: var(--descrip-text, #b8c0cc); + font-weight: 600; + user-select: none; +} + +.ssi-palette[open] .ssi-palette-summary { + margin-bottom: 7px; +} + +.ssi-palette .ssi-hint, +.ssi-palette .ssi-token-group { + margin-top: 6px; +} + +.ssi-hint { + color: var(--descrip-text, #8a93a2); + font-size: 11px; +} + +.ssi-token-group { + display: flex; + flex-wrap: wrap; + align-items: center; + gap: 4px; +} + +.ssi-group-name { + color: #8a93a2; + font-size: 11px; + margin-right: 4px; + min-width: 28px; +} + +.ssi-chip { + padding: 3px 7px; + border: 1px solid var(--border-color, #444); + border-radius: 5px; + background: var(--comfy-input-bg, #2a2a2a); + color: #cdd6e0; + cursor: pointer; + font-size: 11px; + transition: all 0.12s ease; +} + +.ssi-chip:hover { + background: #3a6df0; + color: #fff; + border-color: #3a6df0; +} + +/* 选项网格 */ +.ssi-grid { + display: grid; + grid-template-columns: repeat(3, minmax(0, 1fr)); + gap: 7px; + align-items: end; +} + +.ssi-cell { + display: flex; + flex-direction: column; + gap: 3px; +} + +.ssi-cell-label { + font-size: 11px; + color: var(--descrip-text, #8a93a2); +} + +.ssi-check { + display: flex; + align-items: center; + gap: 6px; + font-size: 12px; + min-height: 30px; +} + +/* 预览区 */ +.ssi-preview { + display: flex; + flex-direction: column; + gap: 5px; + padding: 8px; + border-radius: 6px; + background: rgba(255, 255, 255, 0.035); + border: 1px solid var(--border-color, #444); +} + +.ssi-preview-head { + display: flex; + align-items: center; + justify-content: space-between; +} + +.ssi-preview-title { + font-weight: 600; + color: var(--input-text, #e6e6e6); +} + +.ssi-refresh { + padding: 3px 8px; + border: 1px solid var(--border-color, #555); + border-radius: 5px; + background: transparent; + color: var(--descrip-text, #b8c0cc); + cursor: pointer; +} + +.ssi-refresh:hover { + background: #3a6df0; + color: #fff; +} + +.ssi-path { + font-family: "Consolas", "Menlo", monospace; + overflow-wrap: anywhere; + color: var(--input-text, #e6e6e6); + background: rgba(0, 0, 0, 0.25); + padding: 5px 7px; + border-radius: 4px; + max-height: 38px; + overflow: auto; + font-size: 11px; +} + +.ssi-file { + color: #cdd6e0; + overflow-wrap: anywhere; + font-size: 11px; +} + +.ssi-context { + display: flex; + flex-wrap: wrap; + gap: 6px; +} + +.ssi-context:empty, +.ssi-file:empty { + display: none; +} + +.ssi-ctx-tag { + background: rgba(255, 255, 255, 0.08); + border-radius: 4px; + padding: 2px 8px; + font-size: 11px; + color: #dfe6ee; +} + +.ssi-ctx-tag b { + color: #9fc3ff; + font-weight: 600; +} + +.ssi-status { + font-size: 11px; + color: #8a93a2; +} + +.ssi-status-ok { + color: #6bd08a; +} + +.ssi-status-warn { + color: #f0a35b; +} diff --git a/web/smart_save.js b/web/smart_save.js new file mode 100644 index 0000000..af4ddd7 --- /dev/null +++ b/web/smart_save.js @@ -0,0 +1,445 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const NODE_NAME = "SmartSaveImage"; +const CSS_HREF = "extensions/SmartSaveImage/smart_save.css"; + +// 需要隐藏并由自定义面板托管的原生 widget +const MANAGED_WIDGETS = [ + "root_mode", + "custom_root", + "folder_template", + "filename_template", + "file_format", + "quality", + "collision_mode", + "save_mode", + "manual_model", + "embed_workflow", + "counter_digits", + "png_compression", +]; + +// 占位符说明（分组展示，点击插入到当前聚焦的模板输入框） +const TOKENS = [ + { group: "时间", items: [ + { t: "%date:yyyy-MM-dd%", d: "日期(可自定义格式)" }, + { t: "%year%", d: "年" }, { t: "%month%", d: "月" }, { t: "%day%", d: "日" }, + { t: "%hour%", d: "时" }, { t: "%minute%", d: "分" }, { t: "%second%", d: "秒" }, + ]}, + { group: "模型", items: [ + { t: "%model%", d: "模型名(去扩展名)" }, + { t: "%model_full%", d: "模型完整名" }, + { t: "%unet%", d: "UNet/扩散模型" }, + { t: "%lora%", d: "首个 LoRA" }, + { t: "%vae%", d: "VAE" }, + ]}, + { group: "采样", items: [ + { t: "%seed%", d: "种子" }, { t: "%steps%", d: "步数" }, { t: "%cfg%", d: "CFG" }, + { t: "%sampler%", d: "采样器" }, { t: "%scheduler%", d: "调度器" }, + ]}, + { group: "图片", items: [ + { t: "%width%", d: "宽" }, { t: "%height%", d: "高" }, + { t: "%prompt%", d: "正向提示词" }, { t: "%batch%", d: "批次序号(仅文件名)" }, + ]}, +]; + +function ensureStyles() { + const id = "smart-save-image-css"; + if (document.getElementById(id)) return; + const link = document.createElement("link"); + link.id = id; + link.rel = "stylesheet"; + link.href = CSS_HREF; + document.head.append(link); +} + +function el(tag, className, text) { + const node = document.createElement(tag); + if (className) node.className = className; + if (text != null) node.textContent = text; + return node; +} + +function getWidget(node, name) { + return node.widgets?.find((w) => w.name === name); +} + +function widgetValue(node, name, fallback) { + const w = getWidget(node, name); + return w ? w.value : fallback; +} + +function setWidget(node, name, value) { + const w = getWidget(node, name); + if (!w) return; + w.value = value; + w.callback?.(value); +} + +function hideWidget(node, name) { + const w = getWidget(node, name); + if (!w) return; + w.hidden = true; + w.computeSize = () => [0, -4]; + if (!w.options) w.options = {}; + w.options.serialize = true; +} + +function buildPanel(node) { + let debounceTimer = null; + let previewRequest = 0; + + node.serialize_widgets = true; + for (const name of MANAGED_WIDGETS) hideWidget(node, name); + + const modelWidget = getWidget(node, "manual_model"); + const modelOptions = modelWidget?.options?.values || ["auto"]; + + const root = el("div", "ssi-panel"); + + // ---- 根目录 ---- + const rootRow = el("div", "ssi-field ssi-location"); + rootRow.append(el("label", "ssi-label", "保存位置")); + const rootSeg = el("div", "ssi-seg"); + const ROOTS = [ + { v: "output", label: "输出目录" }, + { v: "custom", label: "自定义" }, + { v: "input", label: "输入" }, + { v: "temp", label: "临时" }, + ]; + const rootButtons = {}; + for (const r of ROOTS) { + const b = el("button", "ssi-seg-btn", r.label); + b.type = "button"; + b.onclick = () => { setWidget(node, "root_mode", r.v); syncRoot(); schedulePreview(); }; + rootButtons[r.v] = b; + rootSeg.append(b); + } + rootRow.append(rootSeg); + + const customRoot = el("input", "ssi-input"); + customRoot.placeholder = "例如 D:\\AI\\output（留空使用输出目录）"; + customRoot.value = widgetValue(node, "custom_root", ""); + customRoot.oninput = () => { setWidget(node, "custom_root", customRoot.value); schedulePreview(); }; + const customRootWrap = el("div", "ssi-field ssi-custom-root"); + customRootWrap.append(el("label", "ssi-cell-label", "自定义路径"), customRoot); + + // ---- 目录模板 ---- + const folderRow = el("div", "ssi-field"); + folderRow.append(el("label", "ssi-label", "子目录规则")); + const folderInput = el("input", "ssi-input ssi-mono"); + folderInput.placeholder = "例如 %date:yyyy-MM-dd%/%model%"; + folderInput.value = widgetValue(node, "folder_template", ""); + folderInput.oninput = () => { setWidget(node, "folder_template", folderInput.value); schedulePreview(); }; + folderInput.onfocus = () => { activeField = folderInput; }; + folderRow.append(folderInput); + + // ---- 文件名模板 ---- + const nameRow = el("div", "ssi-field"); + nameRow.append(el("label", "ssi-label", "文件名规则")); + const nameInput = el("input", "ssi-input ssi-mono"); + nameInput.placeholder = "例如 %model%_%seed%"; + nameInput.value = widgetValue(node, "filename_template", ""); + nameInput.oninput = () => { setWidget(node, "filename_template", nameInput.value); schedulePreview(); }; + nameInput.onfocus = () => { activeField = nameInput; }; + nameRow.append(nameInput); + + let activeField = folderInput; + + // ---- 占位符调色板 ---- + const palette = el("details", "ssi-palette"); + palette.append(el("summary", "ssi-palette-summary", "模板变量")); + palette.append(el("div", "ssi-hint", "先选中目录或文件名输入框，再点击变量插入。")); + for (const grp of TOKENS) { + const gwrap = el("div", "ssi-token-group"); + gwrap.append(el("span", "ssi-group-name", grp.group)); + for (const it of grp.items) { + const chip = el("button", "ssi-chip", it.t.replace(/%/g, "")); + chip.type = "button"; + chip.title = `${it.t} — ${it.d}`; + chip.onclick = () => insertToken(it.t); + gwrap.append(chip); + } + palette.append(gwrap); + } + + function insertToken(token) { + const field = activeField || folderInput; + const start = field.selectionStart ?? field.value.length; + const end = field.selectionEnd ?? field.value.length; + field.value = field.value.slice(0, start) + token + field.value.slice(end); + const pos = start + token.length; + field.setSelectionRange(pos, pos); + field.focus(); + field.dispatchEvent(new Event("input")); + } + + // ---- 模型来源 ---- + const modelRow = el("div", "ssi-row"); + modelRow.append(el("label", "ssi-label", "模型来源")); + const modelSelect = el("select", "ssi-select"); + for (const opt of modelOptions) { + const o = el("option", null, opt === "auto" ? "自动读取工作流" : opt); + o.value = opt; + modelSelect.append(o); + } + modelSelect.value = widgetValue(node, "manual_model", "auto"); + modelSelect.onchange = () => { setWidget(node, "manual_model", modelSelect.value); schedulePreview(); }; + modelRow.append(modelSelect); + + // ---- 格式 / 压缩 / 冲突 / 模式 ---- + const optGrid = el("div", "ssi-grid"); + + const fmtSelect = makeSelect(["png", "jpeg", "webp"], widgetValue(node, "file_format", "png"), + (v) => { setWidget(node, "file_format", v); syncFormat(); schedulePreview(); }); + optGrid.append(labeled("格式", fmtSelect)); + + const compressionInput = el("input", "ssi-input"); + compressionInput.type = "number"; + compressionInput.min = "0"; compressionInput.max = "9"; + compressionInput.value = widgetValue(node, "png_compression", 4); + compressionInput.title = "0 最快且文件最大，9 最慢；所有等级均为无损"; + compressionInput.oninput = () => { + const value = Math.max(0, Math.min(parseInt(compressionInput.value || "0", 10), 9)); + setWidget(node, "png_compression", value); + }; + const compressionWrap = labeled("PNG 压缩", compressionInput); + optGrid.append(compressionWrap); + + const collisionSelect = makeSelect( + [["increment", "自动编号"], ["overwrite", "覆盖"]], + widgetValue(node, "collision_mode", "increment"), + (v) => { setWidget(node, "collision_mode", v); schedulePreview(); }); + optGrid.append(labeled("同名冲突", collisionSelect)); + + const modeSelect = makeSelect( + [["save_and_preview", "保存并预览"], ["save_only", "仅保存"], ["preview_only", "仅预览"]], + widgetValue(node, "save_mode", "save_and_preview"), + (v) => setWidget(node, "save_mode", v)); + optGrid.append(labeled("保存模式", modeSelect)); + + const digitsInput = el("input", "ssi-input"); + digitsInput.type = "number"; + digitsInput.min = "0"; digitsInput.max = "8"; + digitsInput.value = widgetValue(node, "counter_digits", 3); + digitsInput.oninput = () => { setWidget(node, "counter_digits", parseInt(digitsInput.value || "0", 10)); schedulePreview(); }; + optGrid.append(labeled("序号位数", digitsInput)); + + const embedLabel = el("label", "ssi-check"); + const embedBox = el("input"); + embedBox.type = "checkbox"; + embedBox.checked = widgetValue(node, "embed_workflow", true) !== false; + embedBox.onchange = () => setWidget(node, "embed_workflow", embedBox.checked); + embedLabel.append(embedBox, document.createTextNode(" 嵌入工作流")); + optGrid.append(embedLabel); + + // ---- 预览区 ---- + const previewBox = el("div", "ssi-preview"); + const previewHead = el("div", "ssi-preview-head"); + const previewTitle = el("span", "ssi-preview-title", "保存结果预览"); + const refreshBtn = el("button", "ssi-refresh", "刷新"); + refreshBtn.type = "button"; + refreshBtn.onclick = () => runPreview(); + previewHead.append(previewTitle, refreshBtn); + const pathLine = el("div", "ssi-path", "—"); + const fileLine = el("div", "ssi-file", ""); + const ctxLine = el("div", "ssi-context", ""); + const statusLine = el("div", "ssi-status", ""); + previewBox.append(previewHead, pathLine, fileLine, ctxLine, statusLine); + + root.append(rootRow, customRootWrap, folderRow, nameRow, previewBox, palette, modelRow, optGrid); + + function syncFormat() { + compressionWrap.style.display = fmtSelect.value === "png" ? "flex" : "none"; + } + + function syncRoot() { + const mode = widgetValue(node, "root_mode", "output"); + for (const [v, b] of Object.entries(rootButtons)) { + b.classList.toggle("ssi-seg-active", v === mode); + } + customRootWrap.style.display = mode === "custom" ? "flex" : "none"; + } + + function syncFromWidgets() { + customRoot.value = widgetValue(node, "custom_root", ""); + folderInput.value = widgetValue(node, "folder_template", ""); + nameInput.value = widgetValue(node, "filename_template", ""); + const modelValue = widgetValue(node, "manual_model", "auto"); + if (![...modelSelect.options].some((option) => option.value === modelValue)) { + const missingModel = el("option", null, `${modelValue}（当前不可用）`); + missingModel.value = modelValue; + modelSelect.append(missingModel); + } + modelSelect.value = modelValue; + fmtSelect.value = widgetValue(node, "file_format", "png"); + compressionInput.value = widgetValue(node, "png_compression", 4); + collisionSelect.value = widgetValue(node, "collision_mode", "increment"); + modeSelect.value = widgetValue(node, "save_mode", "save_and_preview"); + digitsInput.value = widgetValue(node, "counter_digits", 3); + embedBox.checked = widgetValue(node, "embed_workflow", true) !== false; + syncFormat(); + syncRoot(); + } + + // ---- 调用后端计算预览 ---- + async function runPreview() { + const requestId = ++previewRequest; + statusLine.textContent = "正在计算…"; + statusLine.className = "ssi-status"; + try { + const graph = await app.graphToPrompt(); + const payload = { + prompt: graph?.output || {}, + root_mode: widgetValue(node, "root_mode", "output"), + custom_root: widgetValue(node, "custom_root", ""), + folder_template: widgetValue(node, "folder_template", ""), + filename_template: widgetValue(node, "filename_template", "image"), + file_format: widgetValue(node, "file_format", "png"), + manual_model: widgetValue(node, "manual_model", "auto"), + counter_digits: widgetValue(node, "counter_digits", 3), + collision_mode: widgetValue(node, "collision_mode", "increment"), + batch_size: 1, + }; + const resp = await api.fetchApi("/smartsave/preview", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + const data = await resp.json(); + if (requestId !== previewRequest) return; + if (!data.ok) { + statusLine.textContent = "预览失败：" + (data.error || "未知错误"); + statusLine.className = "ssi-status ssi-status-warn"; + return; + } + pathLine.textContent = data.target; + pathLine.title = data.target; + const examples = (data.example_filenames || []).join(" 、 "); + fileLine.textContent = "示例文件：" + examples; + const c = data.context || {}; + ctxLine.innerHTML = ""; + const chips = [ + ["模型", c.model], ["LoRA", c.lora], ["种子", c.seed], + ["采样器", c.sampler], ["尺寸", c.width && c.height ? `${c.width}x${c.height}` : ""], + ]; + for (const [k, v] of chips) { + if (!v) continue; + const tag = el("span", "ssi-ctx-tag"); + tag.append(el("b", null, k + "："), document.createTextNode(v)); + ctxLine.append(tag); + } + if (data.exists) { + statusLine.textContent = `目录已存在，已有 ${data.existing_count} 张图片`; + statusLine.className = "ssi-status ssi-status-ok"; + } else { + statusLine.textContent = "目录尚不存在，保存时将自动创建"; + statusLine.className = "ssi-status"; + } + } catch (err) { + if (requestId !== previewRequest) return; + statusLine.textContent = "预览异常：" + err; + statusLine.className = "ssi-status ssi-status-warn"; + } + } + + function schedulePreview() { + clearTimeout(debounceTimer); + debounceTimer = setTimeout(runPreview, 350); + } + + node.__ssi_refresh = () => { + syncFromWidgets(); + runPreview(); + }; + + function panelContentHeight() { + const style = getComputedStyle(root); + const visibleChildren = [...root.children].filter((child) => child.offsetHeight > 0); + const gap = parseFloat(style.rowGap || style.gap || "0") || 0; + const padding = (parseFloat(style.paddingTop) || 0) + (parseFloat(style.paddingBottom) || 0); + const childrenHeight = visibleChildren.reduce((sum, child) => sum + child.offsetHeight, 0); + return Math.ceil(padding + childrenHeight + gap * Math.max(visibleChildren.length - 1, 0)); + } + + const widget = node.addDOMWidget("smart_save_panel", "smart-save", root, { + serialize: false, + getMinHeight() { return Math.max(360, panelContentHeight() + 8); }, + getMaxHeight() { return 900; }, + }); + widget.serialize = false; + + function fitNodeToPanel() { + const host = root.parentElement; + if (!host) return; + const desiredHeight = Math.ceil(panelContentHeight() + 50); + const desiredWidth = Math.max(node.size[0], 420); + if (Math.abs(node.size[1] - desiredHeight) > 2 || node.size[0] < 420) { + node.setSize([desiredWidth, desiredHeight]); + app.graph.setDirtyCanvas(true, true); + } + } + + const panelObserver = new ResizeObserver(() => requestAnimationFrame(fitNodeToPanel)); + panelObserver.observe(root); + palette.addEventListener("toggle", () => requestAnimationFrame(fitNodeToPanel)); + + const onRemoved = node.onRemoved; + node.onRemoved = function () { + clearTimeout(debounceTimer); + panelObserver.disconnect(); + onRemoved?.apply(this, arguments); + }; + + syncFromWidgets(); + requestAnimationFrame(() => { + node.setSize([Math.max(node.size[0], 420), Math.max(node.size[1], 500)]); + requestAnimationFrame(fitNodeToPanel); + runPreview(); + }); +} + +// 辅助：带标签的字段 +function labeled(text, control) { + const wrap = el("div", "ssi-cell"); + wrap.append(el("label", "ssi-cell-label", text), control); + return wrap; +} + +// 辅助：下拉框（支持 [value, label] 或纯字符串） +function makeSelect(options, value, onChange) { + const sel = el("select", "ssi-select"); + for (const opt of options) { + const [v, label] = Array.isArray(opt) ? opt : [opt, opt]; + const o = el("option", null, label); + o.value = v; + sel.append(o); + } + sel.value = value; + sel.onchange = () => onChange(sel.value); + return sel; +} + +app.registerExtension({ + name: "Comfy.SmartSaveImage", + init() { + ensureStyles(); + }, + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_NAME) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + buildPanel(this); + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + onConfigure?.apply(this, arguments); + requestAnimationFrame(() => this.__ssi_refresh?.()); + }; + }, +});