From 209683d1de2816f543039792bb83afba216ce6d7 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 20 Apr 2026 16:49:16 +0200 Subject: [PATCH] feat(nodes): add Save Image (Gallery + File Export) node Saves an image to the gallery and additionally exports a copy to the filesystem with a custom filename using prefix and suffix around the gallery UUID. - Filename pattern: {prefix}{uuid}{suffix}.{png|jpg|webp} - UUID matches the gallery entry so exported files can be matched back - Quality input (1-100) for JPG/WEBP, ignored for PNG - Board and Metadata inputs behave like the standard Save Image node - Output directory restricted to subfolders of the InvokeAI outputs folder; absolute paths and path traversal are rejected --- invokeai/app/invocations/image.py | 96 +++++++++ .../invocations/test_save_image_to_file.py | 196 ++++++++++++++++++ 2 files changed, 292 insertions(+) create mode 100644 tests/app/invocations/test_save_image_to_file.py diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 17576a0296..ef00077d9c 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -1,5 +1,6 @@ # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) +from pathlib import Path from typing import Literal, Optional import cv2 @@ -991,6 +992,101 @@ class SaveImageInvocation(BaseInvocation, WithMetadata, WithBoard): return ImageOutput.build(image_dto) +@invocation( + "save_image_to_file", + title="Save Image (Gallery + File Export)", + tags=["image", "export", "file", "save"], + category="image", + version="1.0.0", + use_cache=False, +) +class SaveImageToFileInvocation(BaseInvocation, WithMetadata, WithBoard): + """Saves an image to the gallery (like the standard Save Image node) AND additionally exports a copy + to the filesystem with a custom filename. + + Filename pattern: {prefix}{uuid}{suffix}.{file_format} + - The UUID is the same UUID used for the gallery entry, so the exported file can be matched to the gallery item. + - The gallery entry itself always uses the plain UUID (prefix/suffix apply only to the exported file on disk). + - Board and Metadata inputs behave exactly like the standard Save Image node. + - The export target is restricted to (subfolders of) the InvokeAI outputs folder — absolute paths are rejected. + + Example: prefix="hero_", suffix="_final", file_format="png" → "hero__final.png" + """ + + image: ImageField = InputField(description="The image to save and export") + output_directory: str = InputField( + default="", + description=( + "Target subdirectory (relative to the configured InvokeAI outputs folder) for the exported file. " + "Leave empty to use the outputs folder directly. " + "Example: 'my-exports' → /my-exports/. Nested paths like 'exports/2026' are allowed. " + "Absolute paths and path traversal ('..') are not allowed for security reasons. " + "The directory is created automatically if it doesn't exist." + ), + ) + prefix: str = InputField( + default="", + description="Text prepended to the UUID in the exported filename. Example: 'portrait_' → 'portrait_.png'", + ) + suffix: str = InputField( + default="", + description="Text appended to the UUID (before the extension). Example: '_v2' → '_v2.png'", + ) + file_format: Literal["png", "jpg", "webp"] = InputField( + default="png", + description="File format for the exported file. PNG is lossless; JPG/WEBP are lossy and respect 'quality'.", + ) + quality: int = InputField( + default=95, + ge=1, + le=100, + description="Compression quality for JPG and WEBP (1-100, higher = better quality, larger file). Ignored for PNG.", + ) + + def invoke(self, context: InvocationContext) -> ImageOutput: + image = context.images.get_pil(self.image.image_name) + + image_dto = context.images.save(image=image) + + uuid = Path(image_dto.image_name).stem + + outputs_path = context.config.get().outputs_path + assert outputs_path is not None + + if not self.output_directory: + target_dir = outputs_path + else: + raw = Path(self.output_directory) + if raw.is_absolute() or raw.drive or raw.as_posix().startswith("/"): + raise ValueError( + f"Absolute paths are not allowed in output_directory: {self.output_directory!r}. " + "Use a path relative to the InvokeAI outputs folder." + ) + candidate = (outputs_path / raw).resolve() + outputs_resolved = outputs_path.resolve() + if outputs_resolved != candidate and outputs_resolved not in candidate.parents: + raise ValueError( + f"output_directory must stay within the outputs folder: {self.output_directory!r}" + ) + target_dir = candidate + + target_dir.mkdir(parents=True, exist_ok=True) + + filename = f"{self.prefix}{uuid}{self.suffix}.{self.file_format}" + target_path = target_dir / filename + + if self.file_format == "png": + image.save(target_path, format="PNG") + elif self.file_format == "jpg": + if image.mode in ("RGBA", "LA", "P"): + image = image.convert("RGB") + image.save(target_path, format="JPEG", quality=self.quality) + else: + image.save(target_path, format="WEBP", quality=self.quality) + + return ImageOutput.build(image_dto) + + @invocation( "canvas_paste_back", title="Canvas Paste Back", diff --git a/tests/app/invocations/test_save_image_to_file.py b/tests/app/invocations/test_save_image_to_file.py new file mode 100644 index 0000000000..1c8684e2d1 --- /dev/null +++ b/tests/app/invocations/test_save_image_to_file.py @@ -0,0 +1,196 @@ +"""Tests for SaveImageToFileInvocation.""" + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from PIL import Image + +from invokeai.app.invocations.image import SaveImageToFileInvocation + + +def _make_context(tmp_path: Path, pil_image: Image.Image, gallery_uuid: str = "abc123") -> MagicMock: + context = MagicMock() + context.config.get.return_value.outputs_path = tmp_path + context.images.get_pil.return_value = pil_image + image_dto = MagicMock() + image_dto.image_name = f"{gallery_uuid}.png" + image_dto.width = pil_image.width + image_dto.height = pil_image.height + context.images.save.return_value = image_dto + return context + + +def _build_node(**overrides) -> SaveImageToFileInvocation: + defaults = dict( + id="test", + image={"image_name": "input.png"}, + ) + defaults.update(overrides) + return SaveImageToFileInvocation(**defaults) + + +class TestSaveImageToFileInvocation: + def test_saves_to_gallery(self, tmp_path): + img = Image.new("RGB", (8, 8), (255, 0, 0)) + ctx = _make_context(tmp_path, img) + node = _build_node() + + node.invoke(ctx) + + ctx.images.save.assert_called_once() + assert ctx.images.save.call_args.kwargs["image"] is img + + def test_default_directory_is_outputs_root(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="uuid-1") + node = _build_node() + + node.invoke(ctx) + + assert (tmp_path / "uuid-1.png").exists() + + def test_relative_subdirectory_created(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="uuid-2") + node = _build_node(output_directory="my-exports") + + node.invoke(ctx) + + assert (tmp_path / "my-exports" / "uuid-2.png").exists() + + def test_nested_relative_path(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="uuid-3") + node = _build_node(output_directory="exports/2026/hero") + + node.invoke(ctx) + + assert (tmp_path / "exports" / "2026" / "hero" / "uuid-3.png").exists() + + def test_prefix_and_suffix_applied(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="xyz") + node = _build_node(prefix="hero_", suffix="_final") + + node.invoke(ctx) + + assert (tmp_path / "hero_xyz_final.png").exists() + + def test_prefix_only(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="u") + node = _build_node(prefix="p_") + + node.invoke(ctx) + + assert (tmp_path / "p_u.png").exists() + + def test_suffix_only(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="u") + node = _build_node(suffix="_s") + + node.invoke(ctx) + + assert (tmp_path / "u_s.png").exists() + + def test_filename_uses_gallery_uuid_not_input_uuid(self, tmp_path): + """The exported filename must use the UUID from the new gallery entry, + not the UUID of the input image_name.""" + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="new-uuid") + node = _build_node(image={"image_name": "old-uuid.png"}) + + node.invoke(ctx) + + assert (tmp_path / "new-uuid.png").exists() + assert not (tmp_path / "old-uuid.png").exists() + + @pytest.mark.parametrize( + "bad_path", + [ + "D:/Pictures/Invoke", + "C:/Windows", + "/etc/passwd", + "/tmp/foo", + ], + ) + def test_absolute_paths_rejected(self, tmp_path, bad_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img) + node = _build_node(output_directory=bad_path) + + with pytest.raises(ValueError, match="[Aa]bsolute"): + node.invoke(ctx) + + @pytest.mark.parametrize( + "traversal_path", + [ + "../outside", + "subdir/../../outside", + "..", + ], + ) + def test_path_traversal_rejected(self, tmp_path, traversal_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img) + node = _build_node(output_directory=traversal_path) + + with pytest.raises(ValueError): + node.invoke(ctx) + + def test_png_format(self, tmp_path): + img = Image.new("RGBA", (8, 8), (10, 20, 30, 128)) + ctx = _make_context(tmp_path, img, gallery_uuid="u") + node = _build_node(file_format="png") + + node.invoke(ctx) + + path = tmp_path / "u.png" + assert path.exists() + with Image.open(path) as saved: + assert saved.format == "PNG" + assert saved.mode == "RGBA" + + def test_jpg_format_converts_rgba_to_rgb(self, tmp_path): + img = Image.new("RGBA", (8, 8), (10, 20, 30, 128)) + ctx = _make_context(tmp_path, img, gallery_uuid="u") + node = _build_node(file_format="jpg", quality=80) + + node.invoke(ctx) + + path = tmp_path / "u.jpg" + assert path.exists() + with Image.open(path) as saved: + assert saved.format == "JPEG" + assert saved.mode == "RGB" + + def test_webp_format(self, tmp_path): + img = Image.new("RGB", (8, 8)) + ctx = _make_context(tmp_path, img, gallery_uuid="u") + node = _build_node(file_format="webp", quality=75) + + node.invoke(ctx) + + path = tmp_path / "u.webp" + assert path.exists() + with Image.open(path) as saved: + assert saved.format == "WEBP" + + def test_quality_bounds_enforced_by_pydantic(self): + with pytest.raises(Exception): + _build_node(quality=0) + with pytest.raises(Exception): + _build_node(quality=101) + + def test_output_is_pass_through_of_gallery_dto(self, tmp_path): + img = Image.new("RGB", (16, 24)) + ctx = _make_context(tmp_path, img, gallery_uuid="uuid-out") + node = _build_node() + + result = node.invoke(ctx) + + assert result.image.image_name == "uuid-out.png" + assert result.width == 16 + assert result.height == 24