diff --git a/README.md b/README.md index 09e78e5..1917c79 100644 --- a/README.md +++ b/README.md @@ -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 +- 重新设计的模块化架构 +- 改进的元数据提取功能 +- 更灵活的文件夹管理选项 +- 增强的错误处理和用户反馈 + +## 🤝 支持与反馈 + +如果遇到问题或有改进建议，欢迎反馈！ + +--- + +*享受更智能的图片保存体验！* 🎨 \ No newline at end of file diff --git a/__init__.py b/__init__.py index 14b37da..c14b0ce 100644 --- a/__init__.py +++ b/__init__.py @@ -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 \ No newline at end of file diff --git a/src/SmartSaveImage/__init__.py b/src/SmartSaveImage/__init__.py index e69de29..b132cea 100644 --- a/src/SmartSaveImage/__init__.py +++ b/src/SmartSaveImage/__init__.py @@ -0,0 +1,8 @@ +"""SmartSaveImage 包初始化""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", +] diff --git a/src/SmartSaveImage/core/__init__.py b/src/SmartSaveImage/core/__init__.py new file mode 100644 index 0000000..f179e41 --- /dev/null +++ b/src/SmartSaveImage/core/__init__.py @@ -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", +] diff --git a/src/SmartSaveImage/core/image_utils.py b/src/SmartSaveImage/core/image_utils.py new file mode 100644 index 0000000..f61e661 --- /dev/null +++ b/src/SmartSaveImage/core/image_utils.py @@ -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 diff --git a/src/SmartSaveImage/core/metadata.py b/src/SmartSaveImage/core/metadata.py new file mode 100644 index 0000000..eea3e37 --- /dev/null +++ b/src/SmartSaveImage/core/metadata.py @@ -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) diff --git a/src/SmartSaveImage/core/path_utils.py b/src/SmartSaveImage/core/path_utils.py new file mode 100644 index 0000000..d016de5 --- /dev/null +++ b/src/SmartSaveImage/core/path_utils.py @@ -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}" diff --git a/src/SmartSaveImage/nodes/__init__.py b/src/SmartSaveImage/nodes/__init__.py new file mode 100644 index 0000000..d46c403 --- /dev/null +++ b/src/SmartSaveImage/nodes/__init__.py @@ -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", +] diff --git a/src/SmartSaveImage/nodes/folder_manager.py b/src/SmartSaveImage/nodes/folder_manager.py new file mode 100644 index 0000000..719cbdc --- /dev/null +++ b/src/SmartSaveImage/nodes/folder_manager.py @@ -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 diff --git a/src/SmartSaveImage/nodes/image_saver.py b/src/SmartSaveImage/nodes/image_saver.py new file mode 100644 index 0000000..85444d8 --- /dev/null +++ b/src/SmartSaveImage/nodes/image_saver.py @@ -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 diff --git a/src/SmartSaveImage/utils/__init__.py b/src/SmartSaveImage/utils/__init__.py new file mode 100644 index 0000000..c5f5ea8 --- /dev/null +++ b/src/SmartSaveImage/utils/__init__.py @@ -0,0 +1,7 @@ +"""SmartSaveImage 工具模块""" + +from .validators import InputValidator + +__all__ = [ + "InputValidator", +] diff --git a/src/SmartSaveImage/utils/validators.py b/src/SmartSaveImage/utils/validators.py new file mode 100644 index 0000000..c3bd450 --- /dev/null +++ b/src/SmartSaveImage/utils/validators.py @@ -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 + )