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
This commit is contained in:
Alexander Eichhorn
2026-04-20 16:49:16 +02:00
parent f621bc8fd2
commit 209683d1de
2 changed files with 292 additions and 0 deletions

View File

@@ -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_<uuid>_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' → <outputs>/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_<uuid>.png'",
)
suffix: str = InputField(
default="",
description="Text appended to the UUID (before the extension). Example: '_v2''<uuid>_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",

View File

@@ -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