mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-04-23 03:00:31 -04:00
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:
@@ -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",
|
||||
|
||||
196
tests/app/invocations/test_save_image_to_file.py
Normal file
196
tests/app/invocations/test_save_image_to_file.py
Normal 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
|
||||
Reference in New Issue
Block a user