commit 19cc8a992e00f2dbb68331f19c8ebc70902414a3 Author: kjqwer <2990346238@qq.com> Date: Sat Nov 15 07:42:57 2025 +0800 初始化 diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..d4a2c44 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,21 @@ +# http://editorconfig.org + +root = true + +[*] +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true +charset = utf-8 +end_of_line = lf + +[*.bat] +indent_style = tab +end_of_line = crlf + +[LICENSE] +insert_final_newline = false + +[Makefile] +indent_style = tab diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..0c14fb5 --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,15 @@ +* SmartSaveImage version: +* Python version: +* Operating System: + +### Description + +Describe what you were trying to get done. +Tell us what happened, what went wrong, and what you expected to happen. + +### What I Did + +``` +Paste the command(s) you ran and the output. +If there was a crash, please include the traceback here. +``` diff --git a/.github/workflows/build-pipeline.yml b/.github/workflows/build-pipeline.yml new file mode 100644 index 0000000..43a5a19 --- /dev/null +++ b/.github/workflows/build-pipeline.yml @@ -0,0 +1,33 @@ +name: CI build + +on: + pull_request: + branches: + - master + - main +jobs: + build: + runs-on: ${{ matrix.os }} + env: + PYTHONIOENCODING: "utf8" + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.12"] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[dev] + - name: Run Linting + run: | + ruff check . + - name: Run Tests + run: | + pytest tests/ diff --git a/.github/workflows/publish_node.yml b/.github/workflows/publish_node.yml new file mode 100644 index 0000000..0118542 --- /dev/null +++ b/.github/workflows/publish_node.yml @@ -0,0 +1,21 @@ +name: 📦 Publish to Comfy registry +on: + workflow_dispatch: + push: + tags: + - '*' + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: ♻️ Check out code + uses: actions/checkout@v4 + - name: 📦 Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml new file mode 100644 index 0000000..8d3b5d4 --- /dev/null +++ b/.github/workflows/validate.yml @@ -0,0 +1,13 @@ +name: Validate backwards compatibility + +on: + pull_request: + branches: + - master + - main + +jobs: + validate: + runs-on: ubuntu-latest + steps: + - uses: comfy-org/node-diff@main diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0fedfe7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,114 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# OSX useful to ignore +*.DS_Store +.AppleDouble +.LSOverride + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log + +# Sphinx documentation +docs/_build/ + +# IntelliJ Idea +.idea +*.iml +*.ipr +*.iws + +# PyBuilder +target/ + +# Cookiecutter +output/ +python_boilerplate/ +cookiecutter-pypackage-env/ + +# vscode settings +.history/ +*.code-workspace + +# Frontend extension +node_modules/ +.env +.env.local +.env.development.local +.env.test.local +.env.production.local +npm-debug.log* +yarn-debug.log* +yarn-error.log* +node.zip +.vscode/ +.claude/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..2539924 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,10 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.4.9 + hooks: + # Run the linter. + - id: ruff + args: [ --fix ] + # Run the formatter. + - id: ruff-format diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..534b057 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2025, kj + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..6514960 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,9 @@ +include LICENSE +include README.md + +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif + + diff --git a/README.md b/README.md new file mode 100644 index 0000000..09e78e5 --- /dev/null +++ b/README.md @@ -0,0 +1,67 @@ +# SmartSaveImage + +A node for easy save + +> [!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 + +- A list of features + +## 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 +``` + +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. + +## Publish to Github + +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. + +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). + + +## Tests + +This repo contains unit tests written in Pytest in the `tests/` directory. It is recommended to unit test your custom node. + +- [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 + +## Publishing to Registry + +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. + +- [ ] 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`. + +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! + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..14b37da --- /dev/null +++ b/__init__.py @@ -0,0 +1,442 @@ +"""Top-level package for SmartSaveImage.""" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + +] + +__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", +} + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..efdb854 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,87 @@ +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "SmartSaveImage" +version = "0.0.1" +description = "A node for easy save" +authors = [ + {name = "kj", email = "2990346238@qq.com"} +] +readme = "README.md" +license = {text = "MIT license"} +requires-python = ">=3.10" +classifiers = [] +dependencies = [ + +] + + + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # linting +] + +[project.urls] +Repository = "https://github.com/kjqwer/SmartSaveImage" +BugTracker = "https://github.com/kjqwer/SmartSaveImage/issues" +Documentation = "https://github.com/kjqwer/SmartSaveImage/wiki" + + +[tool.comfy] +PublisherId = "kjqwer" +DisplayName = "SmartSaveImage" +Icon = "" +Tags = [] +Repository = "https://github.com/kjqwer/SmartSaveImage" + +includes = [] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + +[tool.mypy] +files = "." + +# Use strict defaults +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +# Don't require test functions to include types +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +# extend-exclude = ["static", "ci/templates"] +line-length = 140 +src = ["src", "tests"] +target-version = "py39" + +# Add rules to ban exec/eval +[tool.ruff.lint] +select = [ + "S102", # exec-builtin + "S307", # eval-used + "W293", + "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. + # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f +] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" diff --git a/src/SmartSaveImage/__init__.py b/src/SmartSaveImage/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/SmartSaveImage/nodes.py b/src/SmartSaveImage/nodes.py new file mode 100644 index 0000000..fab5cff --- /dev/null +++ b/src/SmartSaveImage/nodes.py @@ -0,0 +1,118 @@ +from inspect import cleandoc +class Example: + """ + A example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "image": ("Image", { "tooltip": "This is an image"}), + "int_field": ("INT", { + "default": 0, + "min": 0, #Minimum value + "max": 4096, #Maximum value + "step": 64, #Slider's step + "display": "number" # Cosmetic only: display as "number" or "slider" + }), + "float_field": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.01, + "round": 0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. + "display": "number"}), + "print_to_screen": (["enable", "disable"],), + "string_field": ("STRING", { + "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node + "default": "Hello World!" + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + #RETURN_NAMES = ("image_output_name",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "test" + + #OUTPUT_NODE = False + #OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node + + CATEGORY = "Example" + + def test(self, image, string_field, int_field, float_field, print_to_screen): + if print_to_screen == "enable": + print(f"""Your input contains: + string_field aka input text: {string_field} + int_field: {int_field} + float_field: {float_field} + """) + #do some processing on the image, in this example I just invert it + image = 1.0 - image + return (image,) + + """ + The node will always be re executed if any of the inputs change but + this method can be used to force the node to execute again even when the inputs don't change. + You can make this node return a number or a string. This value will be compared to the one returned the last time the node was + executed, if it is different the node will be executed again. + This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash + changes between executions the LoadImage node is executed again. + """ + #@classmethod + #def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen): + # return "" + + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "Example": Example +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "Example": "Example Node" +} diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..30c9f69 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Unit test package for SmartSaveImage.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..310609c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,6 @@ +import os +import sys + +# Add the project root directory to Python path +# This allows the tests to import the project +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..95c76f1 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = . # Run tests in the current directory +python_files = test_*.py # Run tests in files that start with "test_" +norecursedirs = .. # Don't run tests in the parent directory diff --git a/tests/test_SmartSaveImage.py b/tests/test_SmartSaveImage.py new file mode 100644 index 0000000..80d6660 --- /dev/null +++ b/tests/test_SmartSaveImage.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python + +"""Tests for `SmartSaveImage` package.""" + +import pytest +from src.SmartSaveImage.nodes import Example + +@pytest.fixture +def example_node(): + """Fixture to create an Example node instance.""" + return Example() + +def test_example_node_initialization(example_node): + """Test that the node can be instantiated.""" + assert isinstance(example_node, Example) + +def test_return_types(): + """Test the node's metadata.""" + assert Example.RETURN_TYPES == ("IMAGE",) + assert Example.FUNCTION == "test" + assert Example.CATEGORY == "Example"