初始化
This commit is contained in:
@@ -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
|
||||
@@ -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.
|
||||
```
|
||||
@@ -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/
|
||||
@@ -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 }}
|
||||
@@ -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
|
||||
+114
@@ -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/
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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!
|
||||
|
||||
+442
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit test package for SmartSaveImage."""
|
||||
@@ -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__), '..')))
|
||||
@@ -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
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user