首次上传
This commit is contained in:
@@ -1,67 +1,185 @@
|
||||
# SmartSaveImage
|
||||
# SmartSaveImage - 智能图片保存节点
|
||||
|
||||
A node for easy save
|
||||
一个功能强大的ComfyUI自定义节点包，提供智能的文件夹管理和图片保存功能。
|
||||
|
||||
> [!NOTE]
|
||||
> This projected was created with a [cookiecutter](https://github.com/Comfy-Org/cookiecutter-comfy-extension) template. It helps you start writing custom nodes without worrying about the Python setup.
|
||||
## 🌟 主要特性
|
||||
|
||||
## Quickstart
|
||||
- **智能文件夹管理** - 自动创建有组织的文件夹结构
|
||||
- **灵活的保存选项** - 支持多种图片格式和质量设置
|
||||
- **元数据嵌入** - 自动提取并保存工作流信息
|
||||
- **批量处理** - 高效处理多张图片
|
||||
- **用户友好** - 直观的界面和丰富的选项
|
||||
|
||||
1. Install [ComfyUI](https://docs.comfy.org/get_started).
|
||||
1. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager)
|
||||
1. Look up this extension in ComfyUI-Manager. If you are installing manually, clone this repository under `ComfyUI/custom_nodes`.
|
||||
1. Restart ComfyUI.
|
||||
## 📦 节点介绍
|
||||
|
||||
# Features
|
||||
### 智能文件夹管理器 (SmartFolderManager)
|
||||
负责创建和管理文件夹结构，从工作流中自动提取元数据。
|
||||
|
||||
- A list of features
|
||||
### 智能图片保存器 (SmartImageSaver)
|
||||
负责保存图片，支持多种格式、压缩选项和元数据嵌入。
|
||||
|
||||
## Develop
|
||||
## 🚀 快速开始
|
||||
|
||||
To install the dev dependencies and pre-commit (will run the ruff hook), do:
|
||||
### 基本使用流程
|
||||
|
||||
```bash
|
||||
cd SmartSaveImage
|
||||
pip install -e .[dev]
|
||||
pre-commit install
|
||||
```
|
||||
1. **添加智能文件夹管理器节点**
|
||||
- 在ComfyUI中搜索"智能文件夹管理器"
|
||||
- 将要保存的图片连接到管理器的images输入
|
||||
|
||||
The `-e` flag above will result in a "live" install, in the sense that any changes you make to your node extension will automatically be picked up the next time you run ComfyUI.
|
||||
2. **配置文件夹结构**
|
||||
- 使用开关控制各层文件夹：日期、模型、种子、提示词、自定义
|
||||
- 设置基础文件夹路径
|
||||
- 可选择连接外部节点（模型、条件、潜在空间）获取更多信息
|
||||
|
||||
## Publish to Github
|
||||
3. **添加智能图片保存器节点**
|
||||
- 搜索"智能图片保存器"
|
||||
- 将文件夹管理器的三个输出全部连接到保存器：
|
||||
- images → images
|
||||
- folder_path → folder_path
|
||||
- metadata_json → metadata_json
|
||||
|
||||
Install Github Desktop or follow these [instructions](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/generating-a-new-ssh-key-and-adding-it-to-the-ssh-agent) for ssh.
|
||||
4. **配置保存选项**
|
||||
- 选择文件格式和质量设置
|
||||
- 设置文件名和预览模式
|
||||
|
||||
1. Create a Github repository that matches the directory name.
|
||||
2. Push the files to Git
|
||||
```
|
||||
git add .
|
||||
git commit -m "project scaffolding"
|
||||
git push
|
||||
```
|
||||
## 📁 文件夹结构配置
|
||||
|
||||
## Writing custom nodes
|
||||
### 灵活的层级控制
|
||||
现在可以通过开关独立控制每一层文件夹的创建：
|
||||
|
||||
An example custom node is located in [node.py](src/SmartSaveImage/nodes.py). To learn more, read the [docs](https://docs.comfy.org/essentials/custom_node_overview).
|
||||
- **日期文件夹** (`enable_date_folder`)
|
||||
- 按日期组织：`2024-11-15/`
|
||||
- 可自定义格式：`yyyy-MM-dd`, `yyyy/MM/dd` 等
|
||||
- 可选择包含时间：`2024-11-15_14-30-25/`
|
||||
|
||||
- **模型文件夹** (`enable_model_folder`)
|
||||
- 按模型组织：`sdxl_base/`
|
||||
- 自动从工作流提取模型名称
|
||||
- 支持手动指定或从模型节点输入
|
||||
|
||||
## Tests
|
||||
- **种子文件夹** (`enable_seed_folder`)
|
||||
- 按种子组织：`seed_12345/`
|
||||
- 自动从工作流提取种子值
|
||||
- 支持手动设置种子
|
||||
|
||||
This repo contains unit tests written in Pytest in the `tests/` directory. It is recommended to unit test your custom node.
|
||||
- **提示词文件夹** (`enable_prompt_folder`)
|
||||
- 按提示词组织：`beautiful_landscape/`
|
||||
- 可设置最大长度，自动清理非法字符
|
||||
- 支持手动输入或从条件节点获取
|
||||
|
||||
- [build-pipeline.yml](.github/workflows/build-pipeline.yml) will run pytest and linter on any open PRs
|
||||
- [validate.yml](.github/workflows/validate.yml) will run [node-diff](https://github.com/Comfy-Org/node-diff) to check for breaking changes
|
||||
- **自定义文件夹** (`enable_custom_folder`)
|
||||
- 完全自定义：`my_project/`
|
||||
- 可以是任意文件夹名称
|
||||
|
||||
## Publishing to Registry
|
||||
### 组合示例
|
||||
- 全开：`2024-11-15/sdxl_base/seed_12345/beautiful_landscape/my_project/`
|
||||
- 仅日期+模型：`2024-11-15/sdxl_base/`
|
||||
- 仅种子+自定义：`seed_12345/experiment_01/`
|
||||
|
||||
If you wish to share this custom node with others in the community, you can publish it to the registry. We've already auto-populated some fields in `pyproject.toml` under `tool.comfy`, but please double-check that they are correct.
|
||||
## 🖼️ 图片保存选项
|
||||
|
||||
You need to make an account on https://registry.comfy.org and create an API key token.
|
||||
### 文件格式支持
|
||||
- **PNG** - 无损压缩，支持透明度
|
||||
- **JPEG** - 有损压缩，文件较小
|
||||
- **WebP** - 现代格式，支持无损和有损
|
||||
- **BMP** - 位图格式
|
||||
- **TIFF** - 高质量格式
|
||||
|
||||
- [ ] Go to the [registry](https://registry.comfy.org). Login and create a publisher id (everything after the `@` sign on your registry profile).
|
||||
- [ ] Add the publisher id into the pyproject.toml file.
|
||||
- [ ] Create an api key on the Registry for publishing from Github. [Instructions](https://docs.comfy.org/registry/publishing#create-an-api-key-for-publishing).
|
||||
- [ ] Add it to your Github Repository Secrets as `REGISTRY_ACCESS_TOKEN`.
|
||||
### 质量设置
|
||||
- **JPEG质量**：1-100（推荐95）
|
||||
- **WebP质量**：1-100（推荐90）
|
||||
- **WebP无损**：启用无损压缩
|
||||
- **PNG压缩**：0-9级别（推荐6）
|
||||
|
||||
A Github action will run on every git push. You can also run the Github action manually. Full instructions [here](https://docs.comfy.org/registry/publishing). Join our [discord](https://discord.com/invite/comfyorg) if you have any questions!
|
||||
### 文件命名选项
|
||||
- **文件名前缀**：自定义前缀
|
||||
- **添加时间戳**：在文件名中包含时间
|
||||
- **添加计数器**：批量保存时的序号
|
||||
- **计数器设置**：起始值和位数
|
||||
|
||||
## 🔧 高级功能
|
||||
|
||||
### 智能元数据获取
|
||||
- **外部节点优先**：连接外部节点时优先从节点获取信息
|
||||
- **工作流自动提取**：没有外部输入时从工作流中自动提取
|
||||
- **图片尺寸检测**：直接从图片数据中获取准确尺寸
|
||||
- **手动补充**：仅在需要时手动输入补充信息
|
||||
|
||||
### 元数据嵌入
|
||||
- **参数记录**：保存采样器、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
|
||||
- 重新设计的模块化架构
|
||||
- 改进的元数据提取功能
|
||||
- 更灵活的文件夹管理选项
|
||||
- 增强的错误处理和用户反馈
|
||||
|
||||
## 🤝 支持与反馈
|
||||
|
||||
如果遇到问题或有改进建议，欢迎反馈！
|
||||
|
||||
---
|
||||
|
||||
*享受更智能的图片保存体验！* 🎨
|
||||
+5
-434
@@ -1,442 +1,13 @@
|
||||
"""Top-level package for SmartSaveImage."""
|
||||
"""智能保存图片 - ComfyUI节点包"""
|
||||
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
|
||||
]
|
||||
|
||||
__author__ = """kj"""
|
||||
__author__ = "kj"
|
||||
__email__ = "2990346238@qq.com"
|
||||
__version__ = "0.0.1"
|
||||
|
||||
from .src.SmartSaveImage.nodes import NODE_CLASS_MAPPINGS
|
||||
from .src.SmartSaveImage.nodes import NODE_DISPLAY_NAME_MAPPINGS
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import hashlib
|
||||
import numpy as np
|
||||
from PIL import Image, PngImagePlugin
|
||||
import piexif
|
||||
import folder_paths
|
||||
import nodes
|
||||
|
||||
class SmartSaveImage:
|
||||
CATEGORY = "IO/Output"
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
FUNCTION = "process"
|
||||
|
||||
token_pattern = re.compile(r"(%[^%]+%)")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"folder_plan": ("STRING", {"default": "", "multiline": True}),
|
||||
"file_format": (["png", "jpeg", "webp"],),
|
||||
"preview_only": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"quality": ("INT", {"default": 100, "min": 1, "max": 100}),
|
||||
"lossless_webp": ("BOOLEAN", {"default": False}),
|
||||
"embed_workflow": ("BOOLEAN", {"default": False}),
|
||||
"add_counter": ("BOOLEAN", {"default": True}),
|
||||
"root_dir": ("STRING", {"default": "output", "multiline": False}),
|
||||
},
|
||||
"hidden": {
|
||||
"id": "UNIQUE_ID",
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
def sanitize(self, s):
|
||||
return re.sub(r"[:*?\"<>|]", "_", s)
|
||||
|
||||
def sanitize_segment(self, s):
|
||||
s = str(s)
|
||||
s = re.sub(r"[:*?\"<>|]", "_", s)
|
||||
s = re.sub(r"\s+", "_", s)
|
||||
s = re.sub(r"[^A-Za-z0-9_\-]", "_", s)
|
||||
return s.strip("_")
|
||||
|
||||
def build_metadata(self, prompt, extra_pnginfo, steps, sampler_name, scheduler, cfg, seed, width, height, modelname):
|
||||
parts = []
|
||||
ptxt = (json.dumps(prompt) if isinstance(prompt, dict) else (prompt or ""))
|
||||
parts.append(str(ptxt).replace("\n", " ").strip())
|
||||
neg = ""
|
||||
if isinstance(extra_pnginfo, dict):
|
||||
neg = extra_pnginfo.get("neg_prompt", "")
|
||||
if neg:
|
||||
parts.append(f"Negative prompt: {str(neg).replace('\n',' ').strip()}")
|
||||
params = []
|
||||
if steps is not None:
|
||||
params.append(f"Steps: {steps}")
|
||||
if sampler_name:
|
||||
if scheduler and scheduler != "normal":
|
||||
params.append(f"Sampler: {sampler_name} {scheduler}")
|
||||
else:
|
||||
params.append(f"Sampler: {sampler_name}")
|
||||
if cfg is not None:
|
||||
params.append(f"CFG Scale: {cfg}")
|
||||
if seed is not None:
|
||||
params.append(f"Seed: {seed}")
|
||||
params.append(f"Size: {width}x{height}")
|
||||
modelhash = None
|
||||
modellabel = None
|
||||
if modelname:
|
||||
try:
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", modelname)
|
||||
h = hashlib.sha256()
|
||||
with open(ckpt_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
h.update(chunk)
|
||||
modelhash = h.hexdigest()[:10]
|
||||
except Exception:
|
||||
modelhash = None
|
||||
modellabel = os.path.splitext(os.path.basename(modelname))[0]
|
||||
if modellabel:
|
||||
if modelhash:
|
||||
params.append(f"Model hash: {modelhash}, Model: {modellabel}")
|
||||
else:
|
||||
params.append(f"Model: {modellabel}")
|
||||
parts.append(", ".join(params))
|
||||
return "\n".join(parts)
|
||||
|
||||
def extract_from_workflow(self, extra_pnginfo):
|
||||
out = {"seed": None, "steps": None, "cfg": None, "sampler_name": None, "scheduler": None, "model": None}
|
||||
wf = None
|
||||
if isinstance(extra_pnginfo, dict):
|
||||
wf = extra_pnginfo.get("workflow")
|
||||
if isinstance(wf, dict):
|
||||
nodes_list = wf.get("nodes") or []
|
||||
for n in nodes_list:
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type") or n.get("type")
|
||||
inputs = n.get("inputs") or {}
|
||||
if ct in ("CheckpointLoaderSimple", "CheckpointLoader", "CheckpointLoaderV2"):
|
||||
m = inputs.get("ckpt_name") or inputs.get("model") or inputs.get("ckpt")
|
||||
if m and not out["model"]:
|
||||
out["model"] = m
|
||||
if ct in ("KSampler", "KSamplerAdvanced"):
|
||||
if inputs.get("seed") is not None:
|
||||
out["seed"] = inputs.get("seed")
|
||||
if inputs.get("steps") is not None:
|
||||
out["steps"] = inputs.get("steps")
|
||||
if inputs.get("cfg") is not None:
|
||||
out["cfg"] = inputs.get("cfg")
|
||||
if inputs.get("sampler_name") is not None:
|
||||
out["sampler_name"] = inputs.get("sampler_name")
|
||||
if inputs.get("scheduler") is not None:
|
||||
out["scheduler"] = inputs.get("scheduler")
|
||||
return out
|
||||
|
||||
def extract_from_prompt(self, prompt):
|
||||
out = {"seed": None, "steps": None, "cfg": None, "sampler_name": None, "scheduler": None, "model": None}
|
||||
if isinstance(prompt, dict):
|
||||
nodes_list = prompt.get("nodes") or []
|
||||
for n in nodes_list:
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type") or n.get("type")
|
||||
inputs = n.get("inputs") or {}
|
||||
if ct in ("CheckpointLoaderSimple", "CheckpointLoader", "CheckpointLoaderV2"):
|
||||
m = inputs.get("ckpt_name") or inputs.get("model") or inputs.get("ckpt")
|
||||
if m and not out["model"]:
|
||||
out["model"] = m
|
||||
if ct in ("KSampler", "KSamplerAdvanced"):
|
||||
if inputs.get("seed") is not None:
|
||||
out["seed"] = inputs.get("seed")
|
||||
if inputs.get("steps") is not None:
|
||||
out["steps"] = inputs.get("steps")
|
||||
if inputs.get("cfg") is not None:
|
||||
out["cfg"] = inputs.get("cfg")
|
||||
if inputs.get("sampler_name") is not None:
|
||||
out["sampler_name"] = inputs.get("sampler_name")
|
||||
if inputs.get("scheduler") is not None:
|
||||
out["scheduler"] = inputs.get("scheduler")
|
||||
return out
|
||||
|
||||
def format_template(self, template, metadata_dict):
|
||||
result = template
|
||||
matches = re.findall(self.token_pattern, template)
|
||||
for seg in matches:
|
||||
inner = seg.strip("%")
|
||||
parts = inner.split(":")
|
||||
key = parts[0]
|
||||
if key == "seed":
|
||||
val = metadata_dict.get("seed")
|
||||
if isinstance(val, (int, float)):
|
||||
if isinstance(val, int) and val < 0:
|
||||
rep = "rand"
|
||||
else:
|
||||
rep = str(val)
|
||||
else:
|
||||
rep = "seed"
|
||||
result = result.replace(seg, self.sanitize_segment(rep))
|
||||
elif key == "width":
|
||||
result = result.replace(seg, self.sanitize_segment(metadata_dict.get("width", "")))
|
||||
elif key == "height":
|
||||
result = result.replace(seg, self.sanitize_segment(metadata_dict.get("height", "")))
|
||||
elif key == "pprompt":
|
||||
raw = metadata_dict.get("prompt", "untitled")
|
||||
txt = str(raw).replace("\n", " ").strip()
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
n = int(parts[1])
|
||||
txt = txt[:n]
|
||||
except Exception:
|
||||
pass
|
||||
result = result.replace(seg, self.sanitize_segment(txt))
|
||||
elif key == "nprompt":
|
||||
raw = metadata_dict.get("negative_prompt", "")
|
||||
txt = str(raw).replace("\n", " ").strip()
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
n = int(parts[1])
|
||||
txt = txt[:n]
|
||||
except Exception:
|
||||
pass
|
||||
result = result.replace(seg, self.sanitize_segment(txt))
|
||||
elif key == "model":
|
||||
m = str(metadata_dict.get("model", "model"))
|
||||
m = os.path.splitext(os.path.basename(m))[0]
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
n = int(parts[1])
|
||||
m = m[:n]
|
||||
except Exception:
|
||||
pass
|
||||
result = result.replace(seg, self.sanitize_segment(m))
|
||||
elif key == "date":
|
||||
from datetime import datetime
|
||||
now = datetime.now()
|
||||
table = {"yyyy": f"{now.year:04d}", "yy": f"{now.year % 100:02d}", "MM": f"{now.month:02d}", "dd": f"{now.day:02d}", "hh": f"{now.hour:02d}", "mm": f"{now.minute:02d}", "ss": f"{now.second:02d}"}
|
||||
fmt = "yyyyMMddhhmmss"
|
||||
if len(parts) >= 2:
|
||||
fmt = parts[1]
|
||||
for k, v in table.items():
|
||||
fmt = fmt.replace(k, v)
|
||||
result = result.replace(seg, fmt)
|
||||
parts = [self.sanitize_segment(p) for p in result.split("/") if p and p.strip()]
|
||||
if not parts:
|
||||
return "ComfyUI"
|
||||
return "/".join(parts)
|
||||
|
||||
def save_batch(self, images, full_output_folder, base_filename, file_format, quality, lossless_webp, embed_workflow, png_parameters_text, extra_pnginfo, add_counter, counter_start):
|
||||
results = []
|
||||
if not os.path.exists(full_output_folder):
|
||||
os.makedirs(full_output_folder, exist_ok=True)
|
||||
for i, image in enumerate(images):
|
||||
arr = 255.0 * image.cpu().numpy()
|
||||
pil = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
||||
fname = base_filename
|
||||
if add_counter:
|
||||
fname += f"_{counter_start + i:05}_"
|
||||
if file_format == "png":
|
||||
file = fname + ".png"
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
if png_parameters_text:
|
||||
pnginfo.add_text("parameters", png_parameters_text)
|
||||
if embed_workflow and extra_pnginfo is not None and isinstance(extra_pnginfo, dict) and "workflow" in extra_pnginfo:
|
||||
pnginfo.add_text("workflow", json.dumps(extra_pnginfo["workflow"]))
|
||||
pil.save(os.path.join(full_output_folder, file), format="PNG", compress_level=4, pnginfo=pnginfo)
|
||||
elif file_format == "jpeg":
|
||||
file = fname + ".jpg"
|
||||
save_kwargs = {"quality": quality, "optimize": True}
|
||||
if png_parameters_text:
|
||||
try:
|
||||
exif_dict = {"Exif": {piexif.ExifIFD.UserComment: b"UNICODE\0" + png_parameters_text.encode("utf-16be")}}
|
||||
exif_bytes = piexif.dump(exif_dict)
|
||||
save_kwargs["exif"] = exif_bytes
|
||||
except Exception:
|
||||
pass
|
||||
pil.save(os.path.join(full_output_folder, file), format="JPEG", **save_kwargs)
|
||||
else:
|
||||
file = fname + ".webp"
|
||||
save_kwargs = {"quality": quality, "lossless": lossless_webp, "method": 0}
|
||||
try:
|
||||
exif_dict = {}
|
||||
if png_parameters_text:
|
||||
exif_dict["Exif"] = {piexif.ExifIFD.UserComment: b"UNICODE\0" + png_parameters_text.encode("utf-16be")}
|
||||
if embed_workflow and extra_pnginfo is not None and isinstance(extra_pnginfo, dict) and "workflow" in extra_pnginfo:
|
||||
exif_dict["0th"] = {piexif.ImageIFD.ImageDescription: "Workflow:" + json.dumps(extra_pnginfo["workflow"])}
|
||||
exif_bytes = piexif.dump(exif_dict)
|
||||
save_kwargs["exif"] = exif_bytes
|
||||
except Exception:
|
||||
pass
|
||||
pil.save(os.path.join(full_output_folder, file), format="WEBP", **save_kwargs)
|
||||
results.append({"filename": file, "subfolder": os.path.basename(os.path.normpath(full_output_folder)), "type": "output"})
|
||||
return results
|
||||
|
||||
def process(self, images, folder_plan, file_format, preview_only, quality=100, lossless_webp=False, embed_workflow=False, add_counter=True, root_dir="", id=None, prompt=None, extra_pnginfo=None):
|
||||
rd = (root_dir or "").strip()
|
||||
base = folder_paths.get_output_directory()
|
||||
if rd.lower() in ("", "output", ".", "./", "/"):
|
||||
output_dir = base
|
||||
elif os.path.isabs(rd):
|
||||
output_dir = rd
|
||||
else:
|
||||
output_dir = os.path.join(base, rd)
|
||||
if not isinstance(images, (list, tuple, np.ndarray)):
|
||||
if len(images.shape) == 3:
|
||||
images = [images]
|
||||
else:
|
||||
images = [img for img in images]
|
||||
h = images[0].shape[0]
|
||||
w = images[0].shape[1]
|
||||
plan = {}
|
||||
try:
|
||||
plan = json.loads(folder_plan) if isinstance(folder_plan, str) and folder_plan.strip() else {}
|
||||
except Exception:
|
||||
plan = {}
|
||||
segments = plan.get("segments") if isinstance(plan, dict) else None
|
||||
meta = plan.get("metadata") if isinstance(plan, dict) else None
|
||||
if not isinstance(meta, dict):
|
||||
ex = self.extract_from_prompt(prompt)
|
||||
if not any(v is not None for v in ex.values()):
|
||||
ex = self.extract_from_workflow(extra_pnginfo)
|
||||
pos_text = None
|
||||
neg_text = None
|
||||
wf = extra_pnginfo.get("workflow") if isinstance(extra_pnginfo, dict) else None
|
||||
if isinstance(wf, dict):
|
||||
for n in wf.get("nodes", []) or []:
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type") or n.get("type")
|
||||
inputs = n.get("inputs") or {}
|
||||
if ct in ("CLIPTextEncode", "CLIPTextEncodeSDXL", "T5TextEncode") and isinstance(inputs.get("text"), str):
|
||||
if pos_text is None:
|
||||
pos_text = inputs.get("text")
|
||||
elif neg_text is None:
|
||||
neg_text = inputs.get("text")
|
||||
if pos_text is None and isinstance(prompt, dict):
|
||||
for n in prompt.get("nodes", []) or []:
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type") or n.get("type")
|
||||
inputs = n.get("inputs") or {}
|
||||
if ct in ("CLIPTextEncode", "CLIPTextEncodeSDXL", "T5TextEncode") and isinstance(inputs.get("text"), str):
|
||||
if pos_text is None:
|
||||
pos_text = inputs.get("text")
|
||||
elif neg_text is None:
|
||||
neg_text = inputs.get("text")
|
||||
meta = {"seed": ex.get("seed"), "steps": ex.get("steps"), "cfg": ex.get("cfg"), "sampler_name": ex.get("sampler_name"), "scheduler": ex.get("scheduler"), "model": ex.get("model"), "width": w, "height": h, "prompt": (pos_text if isinstance(pos_text, str) and pos_text.strip() else (prompt if isinstance(prompt, str) else "untitled")), "negative_prompt": (neg_text if isinstance(neg_text, str) else (extra_pnginfo.get("neg_prompt", "") if isinstance(extra_pnginfo, dict) else ""))}
|
||||
if not isinstance(segments, list) or not segments:
|
||||
template = "ComfyUI/%date:yyyy-MM-dd%/%model%/%seed%/%pprompt:64%"
|
||||
processed_prefix = self.format_template(template, meta)
|
||||
else:
|
||||
cleaned = [self.sanitize_segment(s) for s in segments if isinstance(s, str) and s.strip()]
|
||||
if not cleaned:
|
||||
processed_prefix = "ComfyUI"
|
||||
else:
|
||||
processed_prefix = "/".join(cleaned)
|
||||
if preview_only:
|
||||
res = nodes.PreviewImage().save_images(images, filename_prefix=processed_prefix, prompt=prompt, extra_pnginfo=extra_pnginfo)
|
||||
return {"ui": res.get("ui", {}), "result": (images,)}
|
||||
full_output_folder, base_filename, counter, subfolder, processed_prefix2 = folder_paths.get_save_image_path(processed_prefix, output_dir, w, h)
|
||||
metadata_text = self.build_metadata(meta.get("prompt"), extra_pnginfo, meta.get("steps"), meta.get("sampler_name"), meta.get("scheduler"), meta.get("cfg"), meta.get("seed"), w, h, meta.get("model"))
|
||||
results = self.save_batch(images, full_output_folder, base_filename, file_format, quality, lossless_webp, embed_workflow, metadata_text, extra_pnginfo, add_counter, counter)
|
||||
return {"ui": {"images": results}, "result": (images,)}
|
||||
|
||||
class SmartMetaCollector:
|
||||
CATEGORY = "IO/Output"
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("folder_plan",)
|
||||
FUNCTION = "collect"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
import comfy
|
||||
return {
|
||||
"required": {
|
||||
"mode": (["workflow", "custom"],),
|
||||
"enable_date": ("BOOLEAN", {"default": True}),
|
||||
"date_format": ("STRING", {"default": "yyyy-MM-dd", "multiline": False}),
|
||||
"enable_model": ("BOOLEAN", {"default": True}),
|
||||
"enable_seed": ("BOOLEAN", {"default": True}),
|
||||
"enable_prompt": ("BOOLEAN", {"default": True}),
|
||||
"prompt_len": ("INT", {"default": 64, "min": 1, "max": 512}),
|
||||
},
|
||||
"optional": {
|
||||
"modelname": (folder_paths.get_filename_list("checkpoints"),),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"positive": ("STRING", {"default": "", "multiline": True}),
|
||||
"negative": ("STRING", {"default": "", "multiline": True}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
||||
"width": ("INT", {"default": 0, "min": 0, "max": 16384}),
|
||||
"height": ("INT", {"default": 0, "min": 0, "max": 16384}),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
def collect(self, mode, enable_date, date_format, enable_model, enable_seed, enable_prompt, prompt_len, modelname=None, seed=0, positive="", negative="", sampler_name=None, scheduler=None, width=0, height=0, prompt=None, extra_pnginfo=None):
|
||||
segs = ["ComfyUI"]
|
||||
meta = {"seed": None, "steps": None, "cfg": None, "sampler_name": None, "scheduler": None, "model": None, "width": width if width else None, "height": height if height else None, "prompt": None, "negative_prompt": None}
|
||||
if mode == "workflow":
|
||||
wf = extra_pnginfo.get("workflow") if isinstance(extra_pnginfo, dict) else None
|
||||
ex = SmartSaveImage().extract_from_workflow(extra_pnginfo)
|
||||
meta.update(ex)
|
||||
pos_text = None
|
||||
neg_text = None
|
||||
if isinstance(wf, dict):
|
||||
for n in wf.get("nodes", []) or []:
|
||||
if not isinstance(n, dict):
|
||||
continue
|
||||
ct = n.get("class_type") or n.get("type")
|
||||
inputs = n.get("inputs") or {}
|
||||
if ct in ("CLIPTextEncode", "CLIPTextEncodeSDXL", "T5TextEncode") and isinstance(inputs.get("text"), str):
|
||||
if pos_text is None:
|
||||
pos_text = inputs.get("text")
|
||||
elif neg_text is None:
|
||||
neg_text = inputs.get("text")
|
||||
meta["prompt"] = pos_text if isinstance(pos_text, str) and pos_text.strip() else "untitled"
|
||||
meta["negative_prompt"] = neg_text if isinstance(neg_text, str) else ""
|
||||
else:
|
||||
meta["model"] = modelname or meta["model"]
|
||||
meta["seed"] = seed
|
||||
meta["sampler_name"] = sampler_name
|
||||
meta["scheduler"] = scheduler
|
||||
meta["prompt"] = positive if isinstance(positive, str) and positive.strip() else "untitled"
|
||||
meta["negative_prompt"] = negative if isinstance(negative, str) else ""
|
||||
from datetime import datetime
|
||||
if enable_date:
|
||||
now = datetime.now()
|
||||
table = {"yyyy": f"{now.year:04d}", "yy": f"{now.year % 100:02d}", "MM": f"{now.month:02d}", "dd": f"{now.day:02d}", "hh": f"{now.hour:02d}", "mm": f"{now.minute:02d}", "ss": f"{now.second:02d}"}
|
||||
fmt = date_format or "yyyy-MM-dd"
|
||||
for k, v in table.items():
|
||||
fmt = fmt.replace(k, v)
|
||||
segs.append(SmartSaveImage().sanitize_segment(fmt))
|
||||
if enable_model:
|
||||
m = os.path.splitext(os.path.basename(meta.get("model") or "model"))[0]
|
||||
segs.append(SmartSaveImage().sanitize_segment(m))
|
||||
if enable_seed:
|
||||
s = meta.get("seed")
|
||||
segs.append(SmartSaveImage().sanitize_segment((str(s) if s is not None else "seed")))
|
||||
if enable_prompt:
|
||||
p = str(meta.get("prompt") or "untitled").replace("\n", " ")
|
||||
p = p[:prompt_len] if isinstance(prompt_len, int) and prompt_len > 0 else p
|
||||
segs.append(SmartSaveImage().sanitize_segment(p))
|
||||
plan = {"segments": segs, "metadata": meta}
|
||||
return (json.dumps(plan),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Smart Save Image": SmartSaveImage,
|
||||
"Smart Meta Collector": SmartMetaCollector,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Smart Save Image": "Smart Save Image",
|
||||
"Smart Meta Collector": "Smart Meta Collector",
|
||||
}
|
||||
__version__ = "0.1.0"
|
||||
|
||||
# 从子包导入节点映射
|
||||
from .src.SmartSaveImage import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
@@ -0,0 +1,8 @@
|
||||
"""SmartSaveImage 包初始化"""
|
||||
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
"""SmartSaveImage 核心模块"""
|
||||
|
||||
from .metadata import MetadataExtractor, MetadataBuilder
|
||||
from .path_utils import PathManager
|
||||
from .image_utils import ImageProcessor
|
||||
|
||||
__all__ = [
|
||||
"MetadataExtractor",
|
||||
"MetadataBuilder",
|
||||
"PathManager",
|
||||
"ImageProcessor",
|
||||
]
|
||||
@@ -0,0 +1,168 @@
|
||||
"""图片处理工具模块"""
|
||||
|
||||
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
|
||||
@@ -0,0 +1,201 @@
|
||||
"""元数据提取和构建模块"""
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,203 @@
|
||||
"""路径管理和文件名处理工具"""
|
||||
|
||||
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)
|
||||
|
||||
# 移除或替换非法字符
|
||||
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('_.')
|
||||
|
||||
# 限制长度
|
||||
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:
|
||||
"""解析基础路径"""
|
||||
if not base_folder or base_folder.lower() in ["", "output", "."]:
|
||||
return default_output_dir
|
||||
elif os.path.isabs(base_folder):
|
||||
return base_folder
|
||||
else:
|
||||
return os.path.join(default_output_dir, 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:
|
||||
return os.path.join(base_path, *path_segments)
|
||||
else:
|
||||
return base_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
|
||||
|
||||
# 生成唯一文件名
|
||||
counter = 1
|
||||
while True:
|
||||
new_filename = f"{base_filename}_{counter:03d}{extension}"
|
||||
new_filepath = os.path.join(directory, new_filename)
|
||||
if not os.path.exists(new_filepath):
|
||||
return new_filename
|
||||
counter += 1
|
||||
|
||||
# 防止无限循环
|
||||
if counter > 9999:
|
||||
break
|
||||
|
||||
# 如果还是冲突，使用时间戳
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:17] # 精确到毫秒
|
||||
return f"{base_filename}_{timestamp}{extension}"
|
||||
@@ -0,0 +1,22 @@
|
||||
"""SmartSaveImage 节点模块"""
|
||||
|
||||
from .folder_manager import SmartFolderManager
|
||||
from .image_saver import SmartImageSaver
|
||||
|
||||
# 导出节点映射
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SmartFolderManager": SmartFolderManager,
|
||||
"SmartImageSaver": SmartImageSaver,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SmartFolderManager": "智能文件夹管理",
|
||||
"SmartImageSaver": "智能图片保存",
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
"SmartFolderManager",
|
||||
"SmartImageSaver",
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
]
|
||||
@@ -0,0 +1,401 @@
|
||||
"""智能文件夹管理器节点"""
|
||||
|
||||
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"
|
||||
|
||||
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": "输入图片，用于提取尺寸信息并传递给保存节点"}),
|
||||
"base_folder": ("STRING", {
|
||||
"default": "output",
|
||||
"multiline": False,
|
||||
"tooltip": "基础文件夹路径，可以是相对路径或绝对路径"
|
||||
}),
|
||||
"create_subfolders": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "是否自动创建子文件夹"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
# 文件夹层级开关
|
||||
"enable_date_folder": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "是否创建日期文件夹"
|
||||
}),
|
||||
"enable_model_folder": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "是否创建模型文件夹"
|
||||
}),
|
||||
"enable_seed_folder": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "是否创建种子文件夹"
|
||||
}),
|
||||
"enable_prompt_folder": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "是否创建提示词文件夹"
|
||||
}),
|
||||
"enable_custom_folder": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "是否使用自定义文件夹"
|
||||
}),
|
||||
|
||||
# 日期相关
|
||||
"date_format": ("STRING", {
|
||||
"default": "yyyy-MM-dd",
|
||||
"multiline": False,
|
||||
"tooltip": "日期格式: yyyy年 MM月 dd日 hh时 mm分 ss秒"
|
||||
}),
|
||||
"include_time": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "是否在日期中包含时间"
|
||||
}),
|
||||
|
||||
# 模型信息来源选择
|
||||
"model_source": (["auto", "manual"], {
|
||||
"default": "auto",
|
||||
"tooltip": "模型信息来源：auto=从元数据/外部输入获取，manual=手动选择"
|
||||
}),
|
||||
"manual_model_name": (folder_paths.get_filename_list("checkpoints"), {
|
||||
"tooltip": "手动选择模型（仅在model_source=manual时生效）"
|
||||
}),
|
||||
"model_input": ("MODEL", {
|
||||
"tooltip": "从模型加载器节点输入（仅在model_source=auto时生效）"
|
||||
}),
|
||||
|
||||
# 种子信息
|
||||
"seed_source": (["manual", "external"], {
|
||||
"default": "manual",
|
||||
"tooltip": "种子来源：manual=手动输入，external=外部输入"
|
||||
}),
|
||||
"manual_seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
"tooltip": "手动设置种子值"
|
||||
}),
|
||||
"seed_input": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
"tooltip": "从外部节点输入种子"
|
||||
}),
|
||||
|
||||
# 提示词信息
|
||||
"prompt_source": (["manual", "external"], {
|
||||
"default": "manual",
|
||||
"tooltip": "提示词来源：manual=手动输入，external=外部输入"
|
||||
}),
|
||||
"manual_prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "手动输入正向提示词"
|
||||
}),
|
||||
"manual_negative_prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "手动输入负向提示词"
|
||||
}),
|
||||
"conditioning_positive": ("CONDITIONING", {
|
||||
"tooltip": "从正向条件节点输入（仅在prompt_source=external时生效）"
|
||||
}),
|
||||
"conditioning_negative": ("CONDITIONING", {
|
||||
"tooltip": "从负向条件节点输入（仅在prompt_source=external时生效）"
|
||||
}),
|
||||
|
||||
# 自定义路径
|
||||
"custom_subfolder": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "自定义子文件夹名称"
|
||||
}),
|
||||
|
||||
# 显示选项
|
||||
"model_short_name": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "使用模型短名称（去除扩展名）"
|
||||
}),
|
||||
"prompt_max_length": ("INT", {
|
||||
"default": 50,
|
||||
"min": 10,
|
||||
"max": 200,
|
||||
"tooltip": "提示词文件夹名最大长度"
|
||||
}),
|
||||
"sanitize_names": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "清理文件名中的非法字符"
|
||||
}),
|
||||
},
|
||||
"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=True, 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_source="manual", manual_seed=0, seed_input=0,
|
||||
prompt_source="manual", manual_prompt="", manual_negative_prompt="",
|
||||
conditioning_positive=None, conditioning_negative=None,
|
||||
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"
|
||||
|
||||
# 清理输入字符串
|
||||
manual_prompt = self.validator.sanitize_input_string(manual_prompt, prompt_max_length * 2)
|
||||
manual_negative_prompt = self.validator.sanitize_input_string(manual_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")
|
||||
|
||||
# 种子信息
|
||||
if seed_source == "manual":
|
||||
final_metadata["seed"] = manual_seed
|
||||
else: # external
|
||||
final_metadata["seed"] = seed_input
|
||||
|
||||
# 提示词信息
|
||||
if prompt_source == "manual":
|
||||
final_metadata["positive_prompt"] = manual_prompt if manual_prompt.strip() else None
|
||||
final_metadata["negative_prompt"] = manual_negative_prompt if manual_negative_prompt.strip() else None
|
||||
else: # external
|
||||
# 从conditioning输入提取（这里暂时标记有输入，具体文本仍需要从工作流获取）
|
||||
if conditioning_positive is not None:
|
||||
final_metadata["has_positive_conditioning"] = True
|
||||
if conditioning_negative is not None:
|
||||
final_metadata["has_negative_conditioning"] = True
|
||||
|
||||
# 尺寸信息（直接从图片获取）
|
||||
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
|
||||
|
||||
# 创建文件夹
|
||||
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,
|
||||
"seed_source": seed_source,
|
||||
"prompt_source": prompt_source,
|
||||
"manual_model_name": manual_model_name,
|
||||
"manual_seed": manual_seed,
|
||||
"manual_prompt": manual_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:
|
||||
# 尝试多种方式提取模型名称
|
||||
if hasattr(model_input, 'model_path'):
|
||||
return model_input.model_path
|
||||
elif hasattr(model_input, 'model') and hasattr(model_input.model, 'model_path'):
|
||||
return model_input.model.model_path
|
||||
elif isinstance(model_input, dict):
|
||||
return model_input.get('model_path') or model_input.get('checkpoint_path')
|
||||
except Exception as e:
|
||||
print(f"[SmartFolderManager] 从模型输入提取名称失败: {e}")
|
||||
|
||||
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
|
||||
@@ -0,0 +1,319 @@
|
||||
"""智能图片保存器节点"""
|
||||
|
||||
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"
|
||||
|
||||
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": "要保存的图片（通常来自SmartFolderManager）"}),
|
||||
"folder_path": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "保存路径（来自SmartFolderManager）"
|
||||
}),
|
||||
"metadata_json": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "元数据JSON（来自SmartFolderManager）"
|
||||
}),
|
||||
"filename_prefix": ("STRING", {
|
||||
"default": "image",
|
||||
"multiline": False,
|
||||
"tooltip": "文件名前缀"
|
||||
}),
|
||||
"file_format": (["png", "jpeg", "webp", "bmp", "tiff"], {
|
||||
"default": "png",
|
||||
"tooltip": "图片格式"
|
||||
}),
|
||||
"preview_mode": (["save_and_preview", "preview_only", "save_only"], {
|
||||
"default": "save_and_preview",
|
||||
"tooltip": "保存模式"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
# 文件名选项
|
||||
"add_timestamp": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "在文件名中添加时间戳"
|
||||
}),
|
||||
"add_counter": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "添加计数器（批量保存时）"
|
||||
}),
|
||||
"counter_start": ("INT", {
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 99999,
|
||||
"tooltip": "计数器起始值"
|
||||
}),
|
||||
"counter_padding": ("INT", {
|
||||
"default": 4,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"tooltip": "计数器位数（补零）"
|
||||
}),
|
||||
|
||||
# 图片质量选项
|
||||
"jpeg_quality": ("INT", {
|
||||
"default": 95,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"tooltip": "JPEG质量（1-100）"
|
||||
}),
|
||||
"webp_quality": ("INT", {
|
||||
"default": 90,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"tooltip": "WebP质量（1-100）"
|
||||
}),
|
||||
"webp_lossless": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "WebP无损压缩"
|
||||
}),
|
||||
"png_compression": ("INT", {
|
||||
"default": 6,
|
||||
"min": 0,
|
||||
"max": 9,
|
||||
"tooltip": "PNG压缩级别（0-9）"
|
||||
}),
|
||||
|
||||
# 元数据选项
|
||||
"embed_metadata": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "嵌入元数据到图片"
|
||||
}),
|
||||
"embed_workflow": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "嵌入工作流信息"
|
||||
}),
|
||||
|
||||
# 高级选项
|
||||
"overwrite_existing": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "覆盖已存在的文件"
|
||||
}),
|
||||
"create_backup": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "为覆盖的文件创建备份"
|
||||
}),
|
||||
"optimize_size": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "优化文件大小"
|
||||
}),
|
||||
},
|
||||
"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,)}
|
||||
|
||||
# 确定保存路径
|
||||
if not folder_path:
|
||||
folder_path = folder_paths.get_output_directory()
|
||||
print(f"[SmartImageSaver] 文件夹路径为空，使用默认输出目录: {folder_path}")
|
||||
elif not os.path.exists(folder_path):
|
||||
print(f"[SmartImageSaver] 文件夹不存在，尝试创建: {folder_path}")
|
||||
try:
|
||||
os.makedirs(folder_path, exist_ok=True)
|
||||
except Exception as e:
|
||||
print(f"[SmartImageSaver] 创建文件夹失败: {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 = os.path.join(folder_path, filename)
|
||||
|
||||
# 创建备份
|
||||
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":
|
||||
# 仅保存，不显示预览
|
||||
result["ui"] = {"images": saved_images} if saved_images else {}
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,7 @@
|
||||
"""SmartSaveImage 工具模块"""
|
||||
|
||||
from .validators import InputValidator
|
||||
|
||||
__all__ = [
|
||||
"InputValidator",
|
||||
]
|
||||
@@ -0,0 +1,102 @@
|
||||
"""输入验证工具"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
|
||||
class InputValidator:
|
||||
"""输入验证器"""
|
||||
|
||||
@staticmethod
|
||||
def validate_folder_path(path: str) -> bool:
|
||||
"""验证文件夹路径是否有效"""
|
||||
if not path:
|
||||
return False
|
||||
|
||||
# 检查路径长度
|
||||
if len(path) > 260: # Windows路径长度限制
|
||||
return False
|
||||
|
||||
# 检查非法字符
|
||||
illegal_chars = r'[<>"|?*]'
|
||||
if re.search(illegal_chars, path):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@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
|
||||
)
|
||||
Reference in New Issue
Block a user