1478 lines
59 KiB
Python
1478 lines
59 KiB
Python
import os
|
|
import json
|
|
import pytest
|
|
import platform
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from src.basic_data_handling.path_nodes import (
|
|
PathJoin,
|
|
PathAbspath,
|
|
PathExists,
|
|
PathIsFile,
|
|
PathIsDir,
|
|
PathGetSize,
|
|
PathSplit,
|
|
PathSplitExt,
|
|
PathBasename,
|
|
PathDirname,
|
|
PathGetExtension,
|
|
PathSetExtension,
|
|
PathNormalize,
|
|
PathRelative,
|
|
PathGlob,
|
|
PathExpandVars,
|
|
PathGetCwd,
|
|
PathListDir,
|
|
PathIsAbsolute,
|
|
PathCommonPrefix,
|
|
PathLoadStringFile,
|
|
PathSaveStringFile,
|
|
PathLoadImageRGB,
|
|
PathSaveImageRGB,
|
|
PathLoadImageRGBA,
|
|
PathSaveImageRGBA,
|
|
PathLoadMaskFromAlpha,
|
|
PathLoadMaskFromGreyscale,
|
|
PathInputDir,
|
|
PathOutputDir,
|
|
)
|
|
|
|
|
|
def decode_user_comment(value):
|
|
"""Decode an EXIF UserComment value like piexif/A1111 readers do."""
|
|
if isinstance(value, bytes):
|
|
if value.startswith(b"ASCII\x00\x00\x00"):
|
|
return value[8:].decode("utf-8", errors="ignore")
|
|
if value.startswith(b"UNICODE\x00"):
|
|
return value[8:].decode("utf-16-be", errors="ignore")
|
|
return value.decode("utf-8", errors="ignore")
|
|
return value
|
|
|
|
|
|
def save_success(result):
|
|
"""Success flag of a save-node result.
|
|
|
|
The save nodes return ComfyUI's ``{"ui": ..., "result": ...}`` form so that
|
|
the frontend previews/lists the written image, so the socket outputs live
|
|
under ``result``.
|
|
"""
|
|
assert isinstance(result, dict), result
|
|
return result["result"][0]
|
|
|
|
|
|
def save_passthrough_images(result):
|
|
"""The IMAGE passthrough output of a save-node result."""
|
|
assert isinstance(result, dict), result
|
|
return result["result"][1]
|
|
|
|
|
|
def save_ui_images(result):
|
|
"""The ``ui.images`` entries a save-node result reports for the frontend."""
|
|
assert isinstance(result, dict), result
|
|
return result["ui"]["images"]
|
|
|
|
|
|
def test_path_join():
|
|
node = PathJoin()
|
|
assert node.join_paths("folder", "file.txt") == (os.path.join("folder", "file.txt"),)
|
|
assert node.join_paths("folder", "") == (os.path.join("folder"),)
|
|
assert node.join_paths(".", "") == (".",)
|
|
# Test join with empty first path
|
|
assert node.join_paths("", "file.txt") == ("file.txt",)
|
|
# Test with multiple path components
|
|
nested_path = os.path.join("folder", "subfolder", "file.txt")
|
|
assert node.join_paths("folder", os.path.join("subfolder", "file.txt")) == (nested_path,)
|
|
|
|
|
|
def test_path_load_save_string_file(tmp_path):
|
|
# Test saving a string to a file
|
|
save_node = PathSaveStringFile()
|
|
test_string = "This is a test string\nwith multiple lines"
|
|
file_path = str(tmp_path / "test_string.txt")
|
|
|
|
# Save the string
|
|
assert save_node.save_text(test_string, file_path) == (True,)
|
|
|
|
# Verify file exists
|
|
assert os.path.exists(file_path)
|
|
|
|
# Test loading the string back
|
|
load_node = PathLoadStringFile()
|
|
loaded_string = load_node.load_text(file_path)
|
|
assert loaded_string == (test_string, True)
|
|
|
|
# Test creating directories when saving
|
|
nested_path = str(tmp_path / "nested" / "dir" / "test.txt")
|
|
assert save_node.save_text("Test with nested dirs", nested_path) == (True,)
|
|
assert os.path.exists(nested_path)
|
|
|
|
# Test different encodings
|
|
utf8_text = "UTF-8 text with special chars: 你好, ñ, é, ö"
|
|
utf8_path = str(tmp_path / "utf8_test.txt")
|
|
assert save_node.save_text(utf8_text, utf8_path, encoding="utf-8") == (True,)
|
|
assert load_node.load_text(utf8_path) == (utf8_text, True)
|
|
|
|
# Test error handling
|
|
assert load_node.load_text(str(tmp_path / "nonexistent.txt")) == ("", False)
|
|
|
|
|
|
def test_path_save_string_file_append(tmp_path):
|
|
# Test appending to an existing file
|
|
save_node = PathSaveStringFile()
|
|
load_node = PathLoadStringFile()
|
|
file_path = str(tmp_path / "append_test.txt")
|
|
|
|
# Save initial content
|
|
assert save_node.save_text("First line\n", file_path) == (True,)
|
|
|
|
# Append to the file
|
|
assert save_node.save_text("Second line\n", file_path, append=True) == (True,)
|
|
|
|
# Verify the content was appended, not overwritten
|
|
assert load_node.load_text(file_path) == ("First line\nSecond line\n", True)
|
|
|
|
# Verify default behavior still overwrites
|
|
assert save_node.save_text("Overwritten", file_path) == (True,)
|
|
assert load_node.load_text(file_path) == ("Overwritten", True)
|
|
|
|
|
|
def test_path_abspath():
|
|
node = PathAbspath()
|
|
assert node.get_abspath(".") == (os.path.abspath("."),)
|
|
assert node.get_abspath("folder/file.txt") == (os.path.abspath("folder/file.txt"),)
|
|
# Test with empty string (should return current directory)
|
|
assert node.get_abspath("") == (os.path.abspath(os.getcwd()),)
|
|
# Test with relative path with special characters
|
|
assert node.get_abspath("./file.txt") == (os.path.abspath("./file.txt"),)
|
|
assert node.get_abspath("../file.txt") == (os.path.abspath("../file.txt"),)
|
|
|
|
|
|
def test_path_exists(tmp_path):
|
|
node = PathExists()
|
|
existing_file = tmp_path / "file.txt"
|
|
existing_file.write_text("content")
|
|
existing_dir = tmp_path / "directory"
|
|
existing_dir.mkdir()
|
|
|
|
assert node.check_exists(str(existing_file)) == (True,)
|
|
assert node.check_exists(str(existing_dir)) == (True,)
|
|
assert node.check_exists("nonexistent.file") == (False,)
|
|
# Test with special paths
|
|
assert node.check_exists(".") == (True,)
|
|
# Test with path containing non-existent parent directory
|
|
assert node.check_exists(str(tmp_path / "nonexistent" / "file.txt")) == (False,)
|
|
|
|
|
|
def test_path_is_file(tmp_path):
|
|
node = PathIsFile()
|
|
file = tmp_path / "file.txt"
|
|
file.write_text("content")
|
|
directory = tmp_path / "directory"
|
|
directory.mkdir()
|
|
|
|
assert node.check_is_file(str(file)) == (True,)
|
|
assert node.check_is_file(str(directory)) == (False,)
|
|
assert node.check_is_file("nonexistent.file") == (False,)
|
|
# Test with empty path
|
|
assert node.check_is_file("") == (False,)
|
|
# Test with special paths
|
|
assert node.check_is_file(".") == (False,)
|
|
# Test with symlink to file if platform supports
|
|
if hasattr(os, "symlink"):
|
|
try:
|
|
symlink = tmp_path / "symlink.txt"
|
|
os.symlink(str(file), str(symlink))
|
|
assert node.check_is_file(str(symlink)) == (True,)
|
|
except (OSError, NotImplementedError):
|
|
# Symlinks might not be supported on all platforms
|
|
pass
|
|
|
|
|
|
def test_path_is_dir(tmp_path):
|
|
node = PathIsDir()
|
|
directory = tmp_path / "directory"
|
|
directory.mkdir()
|
|
file = tmp_path / "file.txt"
|
|
file.write_text("content")
|
|
|
|
assert node.check_is_dir(str(directory)) == (True,)
|
|
assert node.check_is_dir(str(file)) == (False,)
|
|
assert node.check_is_dir("nonexistent.dir") == (False,)
|
|
# Test with special paths
|
|
assert node.check_is_dir(".") == (True,)
|
|
assert node.check_is_dir("") == (False,)
|
|
# Test with nested directory
|
|
nested_dir = directory / "nested"
|
|
nested_dir.mkdir()
|
|
assert node.check_is_dir(str(nested_dir)) == (True,)
|
|
# Test with symlink to directory if platform supports
|
|
if hasattr(os, "symlink"):
|
|
try:
|
|
symlink = tmp_path / "symlink_dir"
|
|
os.symlink(str(directory), str(symlink), target_is_directory=True)
|
|
assert node.check_is_dir(str(symlink)) == (True,)
|
|
except (OSError, NotImplementedError, TypeError):
|
|
# Symlinks might not be supported on all platforms
|
|
pass
|
|
|
|
|
|
def test_path_get_size(tmp_path):
|
|
node = PathGetSize()
|
|
file = tmp_path / "file.txt"
|
|
file.write_text("content")
|
|
empty_file = tmp_path / "empty.txt"
|
|
empty_file.write_text("")
|
|
large_file = tmp_path / "large.txt"
|
|
large_file.write_text("a" * 1024) # 1KB file
|
|
|
|
assert node.get_size(str(file)) == (len("content"),)
|
|
assert node.get_size(str(empty_file)) == (0,)
|
|
assert node.get_size(str(large_file)) == (1024,)
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
node.get_size("nonexistent.file")
|
|
|
|
directory = tmp_path / "directory"
|
|
directory.mkdir()
|
|
|
|
with pytest.raises(ValueError):
|
|
node.get_size(str(directory))
|
|
|
|
|
|
def test_path_split():
|
|
node = PathSplit()
|
|
path = os.path.join("folder", "file.txt")
|
|
|
|
assert node.split_path(path) == (os.path.dirname(path), os.path.basename(path))
|
|
assert node.split_path("file.txt") == ("", "file.txt")
|
|
# Test with empty path
|
|
assert node.split_path("") == ("", "")
|
|
# Test with just directory separator
|
|
assert node.split_path("/") == ("/", "")
|
|
# Test with nested paths
|
|
nested_path = os.path.join("folder", "subfolder", "file.txt")
|
|
assert node.split_path(nested_path) == (os.path.dirname(nested_path), "file.txt")
|
|
# Test with trailing separator
|
|
dir_with_sep = "folder" + os.sep
|
|
assert node.split_path(dir_with_sep) == ("folder", "")
|
|
|
|
|
|
def test_path_splitext():
|
|
node = PathSplitExt()
|
|
assert node.split_ext("file.txt") == ("file", ".txt")
|
|
assert node.split_ext("file") == ("file", "")
|
|
# Test with multiple extensions
|
|
assert node.split_ext("file.tar.gz") == ("file.tar", ".gz")
|
|
# Test with path containing directory
|
|
path = os.path.join("folder", "file.txt")
|
|
assert node.split_ext(path) == (os.path.join("folder", "file"), ".txt")
|
|
# Test with hidden file
|
|
assert node.split_ext(".gitignore") == (".gitignore", "")
|
|
# Test with empty string
|
|
assert node.split_ext("") == ("", "")
|
|
|
|
|
|
def test_path_basename():
|
|
node = PathBasename()
|
|
assert node.get_basename("folder/file.txt") == ("file.txt",)
|
|
assert node.get_basename("file.txt") == ("file.txt",)
|
|
# Test with trailing slash
|
|
assert node.get_basename("folder/") == ("",)
|
|
# Test with empty path
|
|
assert node.get_basename("") == ("",)
|
|
# Test with root directory
|
|
assert node.get_basename("/") == ("",)
|
|
# Test with multiple nested directories
|
|
assert node.get_basename("folder/subfolder/file.txt") == ("file.txt",)
|
|
|
|
|
|
def test_path_dirname():
|
|
node = PathDirname()
|
|
assert node.get_dirname("folder/file.txt") == ("folder",)
|
|
assert node.get_dirname("file.txt") == ("",)
|
|
# Test with trailing slash
|
|
assert node.get_dirname("folder/") == ("folder",)
|
|
# Test with empty path
|
|
assert node.get_dirname("") == ("",)
|
|
# Test with root directory
|
|
assert node.get_dirname("/") == ("/",)
|
|
# Test with nested directories
|
|
assert node.get_dirname("folder/subfolder/file.txt") == (os.path.join("folder", "subfolder"),)
|
|
|
|
|
|
def test_path_get_extension():
|
|
node = PathGetExtension()
|
|
assert node.get_extension("file.txt") == (".txt",)
|
|
assert node.get_extension("file") == ("",)
|
|
# Test with multiple dots
|
|
assert node.get_extension("file.tar.gz") == (".gz",)
|
|
# Test with hidden file
|
|
assert node.get_extension(".gitignore") == ("",)
|
|
# Test with path containing directory
|
|
assert node.get_extension("folder/file.txt") == (".txt",)
|
|
# Test with empty string
|
|
assert node.get_extension("") == ("",)
|
|
|
|
|
|
def test_path_set_extension():
|
|
node = PathSetExtension()
|
|
# Test basic extension replacement
|
|
assert node.set_extension("file.txt", ".jpg") == ("file.jpg",)
|
|
# Test adding extension to file without extension
|
|
assert node.set_extension("file", ".png") == ("file.png",)
|
|
# Test with extension without dot prefix
|
|
assert node.set_extension("document.docx", "pdf") == ("document.pdf",)
|
|
# Test with path containing directory
|
|
assert node.set_extension("folder/image.jpg", ".png") == ("folder/image.png",)
|
|
# Test with multiple dots in filename
|
|
assert node.set_extension("archive.tar.gz", ".zip") == ("archive.tar.zip",)
|
|
# Test with empty extension (removes extension)
|
|
assert node.set_extension("config.ini", "") == ("config",)
|
|
# Test with empty string as path
|
|
assert node.set_extension("", ".txt") == (".txt",)
|
|
|
|
|
|
def test_path_load_mask_nodes(tmp_path, monkeypatch):
|
|
# Create test images
|
|
img_size = (64, 64)
|
|
|
|
# Image with alpha channel
|
|
rgba_img = Image.new("RGBA", img_size, color=(255, 0, 0, 128)) # Semi-transparent red
|
|
rgba_path = str(tmp_path / "alpha_mask.png")
|
|
rgba_img.save(rgba_path)
|
|
|
|
# Grayscale image
|
|
gray_img = Image.new("L", img_size)
|
|
# Create a gradient from black to white
|
|
for y in range(img_size[1]):
|
|
for x in range(img_size[0]):
|
|
gray_img.putpixel((x, y), int(255 * (x / img_size[0])))
|
|
gray_path = str(tmp_path / "gray_mask.png")
|
|
gray_img.save(gray_path)
|
|
|
|
# Mock the helper functions
|
|
def mock_load_image_helper(path):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Image file not found: {path}")
|
|
return Image.open(path)
|
|
|
|
def mock_extract_mask_from_alpha(img):
|
|
if "A" in img.getbands():
|
|
alpha = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
|
mask_tensor = 1.0 - torch.from_numpy(alpha)
|
|
return mask_tensor.unsqueeze(0)
|
|
return torch.zeros((1, img.height, img.width), dtype=torch.float32)
|
|
|
|
def mock_extract_mask_from_greyscale(img):
|
|
if img.mode == "L":
|
|
gray = np.array(img).astype(np.float32) / 255.0
|
|
else:
|
|
gray = np.array(img.getchannel("R")).astype(np.float32) / 255.0
|
|
mask_tensor = 1.0 - torch.from_numpy(gray)
|
|
return mask_tensor.unsqueeze(0)
|
|
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_alpha", mock_extract_mask_from_alpha)
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_greyscale", mock_extract_mask_from_greyscale)
|
|
|
|
# Test loading mask from alpha channel
|
|
alpha_node = PathLoadMaskFromAlpha()
|
|
alpha_mask = alpha_node.load_mask_from_alpha(rgba_path)
|
|
|
|
# Verify the mask shape
|
|
assert isinstance(alpha_mask, tuple)
|
|
assert len(alpha_mask) == 2
|
|
assert isinstance(alpha_mask[0], torch.Tensor)
|
|
assert alpha_mask[0].shape == (1, img_size[1], img_size[0])
|
|
|
|
# Test loading mask from grayscale image
|
|
gray_node = PathLoadMaskFromGreyscale()
|
|
gray_mask = gray_node.load_mask_from_greyscale(gray_path)
|
|
|
|
# Verify the mask shape
|
|
assert isinstance(gray_mask, tuple)
|
|
assert len(gray_mask) == 2
|
|
assert isinstance(gray_mask[0], torch.Tensor)
|
|
assert gray_mask[0].shape == (1, img_size[1], img_size[0])
|
|
|
|
# Test with invert option
|
|
inverted_mask = gray_node.load_mask_from_greyscale(gray_path, invert=True)
|
|
assert torch.allclose(inverted_mask[0], 1.0 - gray_mask[0])
|
|
|
|
# Test error handling
|
|
with pytest.raises(FileNotFoundError):
|
|
alpha_node.load_mask_from_alpha(str(tmp_path / "nonexistent.png"))
|
|
|
|
|
|
def test_path_load_save_image_rgba(tmp_path, monkeypatch):
|
|
# Create a test image with transparency
|
|
img_size = (64, 64)
|
|
test_img = Image.new("RGBA", img_size, color=(255, 0, 0, 128)) # Semi-transparent red
|
|
img_path = str(tmp_path / "test_rgba.png")
|
|
test_img.save(img_path)
|
|
|
|
# Mock the load_image_helper and extraction functions
|
|
def mock_load_image_helper(path):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Image file not found: {path}")
|
|
return Image.open(path)
|
|
|
|
def mock_extract_mask_from_alpha(img):
|
|
if "A" in img.getbands():
|
|
alpha = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
|
mask_tensor = 1.0 - torch.from_numpy(alpha)
|
|
return mask_tensor.unsqueeze(0)
|
|
return torch.zeros((1, img.height, img.width), dtype=torch.float32)
|
|
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_alpha", mock_extract_mask_from_alpha)
|
|
|
|
# Test loading an image with alpha
|
|
load_node = PathLoadImageRGBA()
|
|
loaded_img, loaded_mask, success = load_node.load_image_rgba(img_path)
|
|
|
|
# Verify that the returned objects are tensors with the right shapes
|
|
assert isinstance(loaded_img, torch.Tensor)
|
|
assert isinstance(loaded_mask, torch.Tensor)
|
|
assert loaded_img.shape == (1, img_size[1], img_size[0], 3) # (batch, height, width, channels)
|
|
assert loaded_mask.shape == (1, img_size[1], img_size[0]) # (batch, height, width)
|
|
|
|
# Test saving an image with mask
|
|
save_node = PathSaveImageRGBA()
|
|
output_path = str(tmp_path / "output_rgba")
|
|
|
|
# Create a simple red test image tensor and a gradient mask
|
|
red_img = torch.zeros(1, img_size[1], img_size[0], 3)
|
|
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
|
|
|
|
# Create a gradient mask (0.0 to 1.0 from left to right)
|
|
mask = torch.zeros(1, img_size[1], img_size[0])
|
|
for i in range(img_size[0]):
|
|
mask[0, :, i] = i / img_size[0]
|
|
|
|
# Save the image with mask
|
|
assert save_success(save_node.save_image_with_mask(red_img, mask, output_path)) is True
|
|
|
|
# Verify the image was saved with the correct extension
|
|
assert os.path.exists(output_path + ".png")
|
|
|
|
# Test with invert_mask option
|
|
assert save_success(save_node.save_image_with_mask(red_img, mask, str(tmp_path / "inverted"), invert_mask=True)) is True
|
|
assert os.path.exists(str(tmp_path / "inverted.png"))
|
|
|
|
# Test with JPEG format (should switch to PNG for transparency)
|
|
assert save_success(save_node.save_image_with_mask(red_img, mask, str(tmp_path / "jpeg_test"), format="jpg")) is True
|
|
# Should be saved as PNG despite the request for JPEG
|
|
assert os.path.exists(str(tmp_path / "jpeg_test.png"))
|
|
|
|
|
|
def test_path_load_save_image_rgb(tmp_path, monkeypatch):
|
|
# Create a test image
|
|
img_size = (64, 64)
|
|
test_img = Image.new("RGB", img_size, color="red")
|
|
img_path = str(tmp_path / "test_rgb.png")
|
|
test_img.save(img_path)
|
|
|
|
# Mock the load_image_helper function to avoid PIL import issues in testing
|
|
def mock_load_image_helper(path):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Image file not found: {path}")
|
|
return Image.open(path)
|
|
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
|
|
|
|
# Test loading an image
|
|
load_node = PathLoadImageRGB()
|
|
loaded_img = load_node.load_image_rgb(img_path)
|
|
|
|
# Verify that the returned object is a tensor with the right shape
|
|
assert isinstance(loaded_img, tuple)
|
|
assert len(loaded_img) == 2
|
|
assert isinstance(loaded_img[0], torch.Tensor)
|
|
assert loaded_img[0].shape == (1, img_size[1], img_size[0], 3) # (batch, height, width, channels)
|
|
|
|
# Test saving an image
|
|
save_node = PathSaveImageRGB()
|
|
output_path = str(tmp_path / "output_rgb")
|
|
|
|
# Create a simple red test image tensor
|
|
red_img = torch.zeros(1, img_size[1], img_size[0], 3)
|
|
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
|
|
|
|
# Save the image
|
|
assert save_success(save_node.save_image(red_img, output_path)) is True
|
|
|
|
# Verify the image was saved with the correct extension
|
|
assert os.path.exists(output_path + ".png")
|
|
|
|
# Test different formats
|
|
assert save_success(save_node.save_image(red_img, str(tmp_path / "jpeg_test"), format="jpg")) is True
|
|
assert os.path.exists(str(tmp_path / "jpeg_test.jpg"))
|
|
|
|
# Test with directories that don't exist
|
|
nested_path = str(tmp_path / "nested" / "images" / "test_rgb")
|
|
assert save_success(save_node.save_image(red_img, nested_path)) is True
|
|
assert os.path.exists(nested_path + ".png")
|
|
|
|
# Test error handling for loading
|
|
with pytest.raises(FileNotFoundError):
|
|
load_node.load_image_rgb(str(tmp_path / "nonexistent.png"))
|
|
|
|
|
|
def test_path_save_image_rgb_prompt_metadata(tmp_path):
|
|
save_node = PathSaveImageRGB()
|
|
img_size = (16, 16)
|
|
red_img = torch.zeros(1, img_size[1], img_size[0], 3)
|
|
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
|
|
|
|
# No metadata when neither prompt nor negative prompt is provided
|
|
plain_path = str(tmp_path / "plain")
|
|
assert save_success(save_node.save_image(red_img, plain_path)) is True
|
|
with Image.open(plain_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") is None
|
|
|
|
# Both prompt and negative prompt provided
|
|
full_path = str(tmp_path / "full")
|
|
assert save_success(save_node.save_image(red_img, full_path, prompt="a red square", negative_prompt="blurry, low quality")) is True
|
|
with Image.open(full_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality"
|
|
|
|
# Only the prompt is provided
|
|
pos_path = str(tmp_path / "pos_only")
|
|
assert save_success(save_node.save_image(red_img, pos_path, prompt="only positive")) is True
|
|
with Image.open(pos_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") == "only positive"
|
|
|
|
# Only the negative prompt is provided
|
|
neg_path = str(tmp_path / "neg_only")
|
|
assert save_success(save_node.save_image(red_img, neg_path, negative_prompt="only negative")) is True
|
|
with Image.open(neg_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") == "Negative prompt: only negative"
|
|
|
|
# JPEG embeds the prompt into the EXIF UserComment and ImageDescription fields
|
|
jpg_path = str(tmp_path / "jpeg_meta")
|
|
assert (
|
|
save_success(save_node.save_image(red_img, jpg_path, format="jpg", prompt="a red square", negative_prompt="blurry, low quality"))
|
|
is True
|
|
)
|
|
assert os.path.exists(jpg_path + ".jpg")
|
|
with Image.open(jpg_path + ".jpg") as img:
|
|
img.load()
|
|
exif = img.getexif()
|
|
expected = "a red square\nNegative prompt: blurry, low quality"
|
|
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected
|
|
assert exif.get(0x010E) == expected
|
|
|
|
# WEBP embeds the prompt into the EXIF UserComment field (no ImageDescription)
|
|
webp_path = str(tmp_path / "webp_meta")
|
|
assert (
|
|
save_success(save_node.save_image(red_img, webp_path, format="webp", prompt="a red square", negative_prompt="blurry, low quality"))
|
|
is True
|
|
)
|
|
with Image.open(webp_path + ".webp") as img:
|
|
img.load()
|
|
exif_bytes = img.info.get("exif", b"")
|
|
if exif_bytes.startswith(b"Exif\x00\x00"):
|
|
exif_bytes = exif_bytes[6:]
|
|
exif = Image.Exif()
|
|
exif.load(exif_bytes)
|
|
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected
|
|
assert 0x010E not in exif
|
|
|
|
|
|
def test_path_save_image_rgba_prompt_metadata(tmp_path):
|
|
save_node = PathSaveImageRGBA()
|
|
img_size = (16, 16)
|
|
red_img = torch.zeros(1, img_size[1], img_size[0], 3)
|
|
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
|
|
mask = torch.zeros(1, img_size[1], img_size[0])
|
|
|
|
# No metadata when neither prompt nor negative prompt is provided
|
|
plain_path = str(tmp_path / "plain")
|
|
assert save_success(save_node.save_image_with_mask(red_img, mask, plain_path)) is True
|
|
with Image.open(plain_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") is None
|
|
|
|
# Both prompt and negative prompt provided
|
|
full_path = str(tmp_path / "full")
|
|
assert (
|
|
save_success(save_node.save_image_with_mask(red_img, mask, full_path, prompt="a red square", negative_prompt="blurry, low quality"))
|
|
is True
|
|
)
|
|
with Image.open(full_path + ".png") as img:
|
|
img.load()
|
|
assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality"
|
|
|
|
|
|
def test_path_save_image_jxl_prompt_metadata(tmp_path):
|
|
pytest.importorskip("pillow_jxl")
|
|
save_node = PathSaveImageRGB()
|
|
img_size = (16, 16)
|
|
red_img = torch.zeros(1, img_size[1], img_size[0], 3)
|
|
red_img[0, :, :, 0] = 1.0 # Red channel set to 1
|
|
|
|
prompt = "a red square <lora:x:1.0>"
|
|
negative_prompt = "blurry, low quality"
|
|
jxl_path = str(tmp_path / "jxl_meta")
|
|
assert save_success(save_node.save_image(red_img, jxl_path, format="jxl", prompt=prompt, negative_prompt=negative_prompt)) is True
|
|
assert os.path.exists(jxl_path + ".jxl")
|
|
|
|
# JXL embeds the prompt into the EXIF UserComment field
|
|
with Image.open(jxl_path + ".jxl") as img:
|
|
img.load()
|
|
exif_bytes = img.info.get("exif", b"")
|
|
if exif_bytes.startswith(b"Exif\x00\x00"):
|
|
exif_bytes = exif_bytes[6:]
|
|
exif = Image.Exif()
|
|
exif.load(exif_bytes)
|
|
expected = "a red square <lora:x:1.0>\nNegative prompt: blurry, low quality"
|
|
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected
|
|
|
|
# JXL also embeds the prompt (XML-escaped) in an XMP xml box
|
|
raw = (tmp_path / "jxl_meta.jxl").read_bytes()
|
|
assert b"xml " in raw
|
|
assert b"<dc:description>" in raw
|
|
escaped = "a red square <lora:x:1.0>\nNegative prompt: blurry, low quality"
|
|
assert escaped.encode("utf-8") in raw
|
|
|
|
|
|
def _batch_rgb(n, size=(16, 16)):
|
|
"""Build an n-frame RGB tensor, frame i filled with a distinct colour."""
|
|
frames = torch.zeros(n, size[1], size[0], 3)
|
|
for i in range(n):
|
|
frames[i, :, :, i % 3] = (i + 1) / float(n)
|
|
return frames
|
|
|
|
|
|
def test_path_save_image_format_is_dropdown():
|
|
"""The format input is a combo/dropdown whose entries are derived from what
|
|
the installed Pillow/plugins can actually save (not a hard-coded constant)."""
|
|
import src.basic_data_handling.path_nodes as pn
|
|
|
|
# Must include the classic formats and be reflect the live library/plugins.
|
|
assert pn._IMAGE_SAVE_FORMATS
|
|
assert pn._IMAGE_SAVE_FORMATS[0] == "png"
|
|
|
|
# png & alpha-capable core are always surfaced on the IMAGE+MASK node.
|
|
assert pn._IMAGE_SAVE_ALPHA_FORMATS
|
|
assert pn._IMAGE_SAVE_ALPHA_FORMATS[0] == "png"
|
|
assert "png" in pn._IMAGE_SAVE_ALPHA_FORMATS
|
|
|
|
rgb_types = PathSaveImageRGB.INPUT_TYPES()
|
|
formats = rgb_types["optional"]["format"]
|
|
assert isinstance(formats[0], list) and formats[0] == pn._IMAGE_SAVE_FORMATS
|
|
assert formats[1]["default"] == "png"
|
|
|
|
rgba_types = PathSaveImageRGBA.INPUT_TYPES()
|
|
formats_rgba = rgba_types["optional"]["format"]
|
|
# the IMAGE+MASK node offers only alpha-capable formats (jpg is absent, png
|
|
# present) but its runtime value is still a STRING-typed text.
|
|
assert isinstance(formats_rgba[0], list)
|
|
assert formats_rgba[0] == pn._IMAGE_SAVE_ALPHA_FORMATS
|
|
assert "jpg" not in formats_rgba[0]
|
|
assert formats_rgba[1]["default"] == "png"
|
|
|
|
# Downstream aliases still accepted: jpg -> canonical jpg token.
|
|
assert pn._normalize_image_format("jpeg") == "jpg"
|
|
|
|
# The format is still text at runtime, so it can be driven by a STRING link.
|
|
for f in formats[0]:
|
|
assert isinstance(f, str)
|
|
|
|
# Both nodes expose a use_prefix_mode toggle.
|
|
assert rgba_types["optional"]["use_prefix_mode"][1]["default"] is False
|
|
|
|
|
|
def test_path_save_image_formats_are_library_derived():
|
|
"""The dropdown is derived from the live Pillow/plugin library, so it must
|
|
contain only formats that really save, exclude non-photographic ones, and
|
|
grow/shrink with the environment."""
|
|
import src.basic_data_handling.path_nodes as pn
|
|
|
|
# Non-photographic raster ids are never surfaced.
|
|
assert not any(t in pn._NON_PHOTOGRAPHIC_RASTER_IDS for t in map(str.upper, pn._IMAGE_SAVE_FORMATS))
|
|
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
for fmt in pn._IMAGE_SAVE_FORMATS:
|
|
out = f"/tmp/_bdh_fmt_{fmt}"
|
|
# RGB node writes any offered format
|
|
assert save_success(save_node.save_image(img, out, format=fmt)) is True, fmt
|
|
# correct file extension appended
|
|
expected = f"{out}.{fmt}"
|
|
assert os.path.exists(expected), fmt
|
|
os.remove(expected)
|
|
|
|
# Some classic formats are always present when Pillow can write them.
|
|
for fmt in ("png", "jpg", "webp"):
|
|
assert fmt in pn._IMAGE_SAVE_FORMATS
|
|
|
|
|
|
def test_path_save_image_rgb_multi_frame_suffix(tmp_path):
|
|
save_node = PathSaveImageRGB()
|
|
frames = _batch_rgb(3)
|
|
out = str(tmp_path / "frames")
|
|
|
|
assert save_success(save_node.save_image(frames, out)) is True
|
|
expected = [
|
|
str(tmp_path / "frames_00000.png"),
|
|
str(tmp_path / "frames_00001.png"),
|
|
str(tmp_path / "frames_00002.png"),
|
|
]
|
|
for path in expected:
|
|
assert os.path.exists(path)
|
|
|
|
|
|
def test_path_save_image_prefix_mode_mirrors_comfy(tmp_path, monkeypatch):
|
|
"""use_prefix_mode writes into the output dir with ComfyUI-style
|
|
``prefix_00001_.png`` naming and auto-increments on repeated saves."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
prefix = "my_subdir/shot"
|
|
|
|
assert save_success(save_node.save_image(img, prefix, use_prefix_mode=True)) is True
|
|
first = tmp_path / "my_subdir" / "shot_00001_.png"
|
|
assert first.exists()
|
|
|
|
# A second save must not overwrite the first file.
|
|
assert save_success(save_node.save_image(img, prefix, use_prefix_mode=True)) is True
|
|
assert (tmp_path / "my_subdir" / "shot_00002_.png").exists()
|
|
|
|
|
|
def test_path_save_image_prefix_mode_batch(tmp_path, monkeypatch):
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
frames = _batch_rgb(2)
|
|
|
|
assert save_success(save_node.save_image(frames, "batch", use_prefix_mode=True)) is True
|
|
assert (tmp_path / "batch_00001_.png").exists()
|
|
assert (tmp_path / "batch_00002_.png").exists()
|
|
|
|
|
|
def test_path_save_image_prefix_mode_template_tokens(tmp_path, monkeypatch):
|
|
"""Int/width/height tokens are substituted like ComfyUI's backend before use."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1, size=(8, 4)) # width=8, height=4
|
|
|
|
assert save_success(save_node.save_image(img, "w%width%_h%height%", use_prefix_mode=True)) is True
|
|
assert (tmp_path / "w8_h4_00001_.png").exists()
|
|
|
|
|
|
def test_path_save_image_prefix_mode_date_token(tmp_path, monkeypatch):
|
|
"""``%date:yyyy-MM-dd%`` is expanded to the real date folder (not a literal
|
|
``%date:...%`` directory) in prefix mode."""
|
|
from datetime import datetime
|
|
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
|
|
assert save_success(save_node.save_image(img, "%date:yyyy-MM-dd%/dated_test", use_prefix_mode=True)) is True
|
|
# The literal token must never be used as a folder name.
|
|
assert not (tmp_path / "%date:yyyy-MM-dd%").exists()
|
|
# A folder named after today's date holds the numbered file.
|
|
today = datetime.now().strftime("%Y-%m-%d")
|
|
assert (tmp_path / today / "dated_test_00001_.png").exists()
|
|
|
|
|
|
def test_path_save_image_path_mode_date_token(tmp_path):
|
|
"""``%date:...%`` templates also work in the plain path mode."""
|
|
from datetime import datetime
|
|
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
today = datetime.now().strftime("%Y-%m-%d")
|
|
out = os.path.join(str(tmp_path), "%date:yyyy-MM-dd%", "dated")
|
|
|
|
assert save_success(save_node.save_image(img, out)) is True
|
|
assert not (tmp_path / "%date:yyyy-MM-dd%").exists()
|
|
assert (tmp_path / today / "dated.png").exists()
|
|
|
|
|
|
def test_path_save_image_empty_path_returns_false():
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
result = save_node.save_image(img, "")
|
|
assert save_success(result) is False
|
|
# The images are passed through even when nothing was written.
|
|
assert save_passthrough_images(result) is img
|
|
assert save_ui_images(result) == []
|
|
|
|
|
|
def test_path_save_image_rgba_prefix_mode_and_transparency(tmp_path, monkeypatch):
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGBA()
|
|
size = (8, 8)
|
|
img = torch.zeros(1, size[1], size[0], 3)
|
|
img[0, :, :, 0] = 1.0
|
|
mask = torch.zeros(1, size[1], size[0])
|
|
|
|
assert save_success(save_node.save_image_with_mask(img, mask, "alpha", use_prefix_mode=True)) is True
|
|
path = tmp_path / "alpha_00001_.png"
|
|
assert path.exists()
|
|
|
|
# The right half is masked (alpha 255/opaque where mask==0); check alpha via RGBA.
|
|
with Image.open(path) as loaded:
|
|
loaded.load()
|
|
# PNG metadata must be empty (no prompt given).
|
|
assert loaded.text.get("parameters") is None
|
|
|
|
|
|
def test_path_save_image_rgba_multi_frame(tmp_path):
|
|
save_node = PathSaveImageRGBA()
|
|
size = (8, 8)
|
|
frames = torch.zeros(2, size[1], size[0], 3)
|
|
frames[0, :, :, 0] = 1.0
|
|
frames[1, :, :, 1] = 1.0
|
|
mask = torch.zeros(1, size[1], size[0]) # single mask shared across frames
|
|
|
|
assert save_success(save_node.save_image_with_mask(frames, mask, str(tmp_path / "rgba_batch"))) is True
|
|
assert (tmp_path / "rgba_batch_00000.png").exists()
|
|
assert (tmp_path / "rgba_batch_00001.png").exists()
|
|
|
|
|
|
def test_path_save_image_reports_preview_entries(tmp_path, monkeypatch):
|
|
"""Files saved into ComfyUI's output folder are reported as ``ui.images``,
|
|
which is what makes the frontend preview them and list them with the
|
|
generated images (same payload as the built-in "Save Image" node)."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
frames = _batch_rgb(2)
|
|
|
|
result = save_node.save_image(frames, "my_subdir/shot", use_prefix_mode=True)
|
|
|
|
assert save_ui_images(result) == [
|
|
{"filename": "shot_00001_.png", "subfolder": "my_subdir", "type": "output"},
|
|
{"filename": "shot_00002_.png", "subfolder": "my_subdir", "type": "output"},
|
|
]
|
|
# The images are passed through unchanged, like the built-in node does.
|
|
assert save_passthrough_images(result) is frames
|
|
|
|
|
|
def test_path_save_image_rgba_reports_preview_entries(tmp_path, monkeypatch):
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGBA()
|
|
size = (8, 8)
|
|
img = torch.zeros(1, size[1], size[0], 3)
|
|
mask = torch.zeros(1, size[1], size[0])
|
|
|
|
result = save_node.save_image_with_mask(img, mask, "alpha", use_prefix_mode=True)
|
|
|
|
assert save_ui_images(result) == [{"filename": "alpha_00001_.png", "subfolder": "", "type": "output"}]
|
|
|
|
|
|
def test_path_save_image_plain_path_inside_output_dir_is_previewed(tmp_path, monkeypatch):
|
|
"""Plain path mode is previewable as long as the file lands in the output folder."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
|
|
result = save_node.save_image(img, str(tmp_path / "sub" / "plain"))
|
|
|
|
assert save_ui_images(result) == [{"filename": "plain.png", "subfolder": "sub", "type": "output"}]
|
|
|
|
|
|
def test_path_save_image_temp_dir_is_previewed_as_temp(tmp_path, monkeypatch):
|
|
"""Files in ComfyUI's temp folder are reported with ``type: "temp"``."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path / "out"))
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_temp_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
|
|
result = save_node.save_image(img, str(tmp_path / "scratch"))
|
|
|
|
assert save_ui_images(result) == [{"filename": "scratch.png", "subfolder": "", "type": "temp"}]
|
|
|
|
|
|
def test_path_save_image_outside_comfy_folders_has_no_preview(tmp_path, monkeypatch):
|
|
"""The file is still written, but nothing is reported for it: ComfyUI could
|
|
not serve it through /view, so a preview entry would be a broken image."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path / "out"))
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_temp_directory", lambda: str(tmp_path / "tmp"))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
out = str(tmp_path / "elsewhere" / "image")
|
|
|
|
result = save_node.save_image(img, out)
|
|
|
|
assert save_success(result) is True
|
|
assert save_ui_images(result) == []
|
|
assert os.path.exists(out + ".png")
|
|
|
|
|
|
def test_path_save_image_exposes_workflow_inputs_and_image_passthrough():
|
|
"""Drop-in superset of the built-in node: it takes ComfyUI's hidden workflow
|
|
prompt info and passes the images through as a second output."""
|
|
for cls in (PathSaveImageRGB, PathSaveImageRGBA):
|
|
assert cls.INPUT_TYPES()["hidden"] == {"workflow_prompt": "PROMPT", "workflow_extra_pnginfo": "EXTRA_PNGINFO"}
|
|
assert cls.RETURN_TYPES == ("BOOLEAN", "IMAGE")
|
|
assert cls.RETURN_NAMES == ("success", "images")
|
|
assert cls.OUTPUT_NODE is True
|
|
|
|
|
|
def test_path_save_image_embeds_workflow_and_explicit_prompt(tmp_path, monkeypatch):
|
|
"""Both metadata channels coexist in a PNG: the workflow (all prompts of the
|
|
graph, needed so dragging the file back restores the workflow) and the
|
|
explicit prompt of this save as ``parameters`` (easy to extract, shown by
|
|
image viewers)."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
frames = _batch_rgb(2)
|
|
workflow_prompt = {
|
|
"3": {"class_type": "KSampler", "inputs": {"seed": 1}},
|
|
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "prompt of a different node"}},
|
|
}
|
|
extra_pnginfo = {"workflow": {"nodes": [{"id": 3}, {"id": 6}]}}
|
|
|
|
result = save_node.save_image(
|
|
frames,
|
|
"meta",
|
|
use_prefix_mode=True,
|
|
prompt="the prompt of this image",
|
|
negative_prompt="bad",
|
|
workflow_prompt=workflow_prompt,
|
|
workflow_extra_pnginfo=extra_pnginfo,
|
|
)
|
|
assert save_success(result) is True
|
|
|
|
# Every frame of the batch carries both channels.
|
|
for name in ("meta_00001_.png", "meta_00002_.png"):
|
|
with Image.open(tmp_path / name) as img:
|
|
img.load()
|
|
assert img.text["parameters"] == "the prompt of this image\nNegative prompt: bad"
|
|
assert json.loads(img.text["prompt"]) == workflow_prompt
|
|
assert json.loads(img.text["workflow"]) == extra_pnginfo["workflow"]
|
|
|
|
|
|
#: The sampler settings an A1111 infotext reader knows. Readers (Civitai and
|
|
#: similar) scan the ``parameters`` text from its end for a line holding at least
|
|
#: three of them - ``Sampler`` among them - and discard the whole text when there
|
|
#: is no such line, so the prompt would not be extracted either.
|
|
_A1111_SAMPLER_KEYS = ("Sampler", "CFG scale", "Seed", "Steps", "ENSD", "Schedule type")
|
|
|
|
|
|
def read_as_a1111_parameters(parameters_text):
|
|
"""Return the sampler settings a reader finds in a ``parameters`` text."""
|
|
for line in reversed(parameters_text.splitlines()):
|
|
settings = {}
|
|
for item in line.split(","):
|
|
key, separator, value = item.partition(":")
|
|
if separator and key.strip() in _A1111_SAMPLER_KEYS:
|
|
settings[key.strip()] = value.strip()
|
|
if len(settings) >= 3 and "Sampler" in settings:
|
|
return settings
|
|
return None
|
|
|
|
|
|
def _sampler_workflow(**overrides):
|
|
"""An API workflow prompt holding one KSampler with literal settings."""
|
|
settings = {"steps": 8, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "seed": 51, "denoise": 1.0}
|
|
settings.update(overrides)
|
|
return {"3": {"class_type": "KSampler", "inputs": {"model": ["10", 0], "positive": ["6", 0], **settings}}}
|
|
|
|
|
|
def test_path_save_image_parameters_text_carries_the_sampler_settings(tmp_path):
|
|
"""The explicit prompt is only extracted by readers like Civitai when the
|
|
``parameters`` text ends with a ``Steps: ...`` line, so the save node adds the
|
|
sampler settings of the workflow below the prompt."""
|
|
save_node = PathSaveImageRGB()
|
|
out = str(tmp_path / "a1111_meta")
|
|
|
|
result = save_node.save_image(
|
|
_batch_rgb(1),
|
|
out,
|
|
prompt="the prompt of this image",
|
|
negative_prompt="bad",
|
|
workflow_prompt=_sampler_workflow(),
|
|
workflow_extra_pnginfo={"workflow": {"nodes": [{"id": 3}]}},
|
|
)
|
|
assert save_success(result) is True
|
|
|
|
with Image.open(out + ".png") as img:
|
|
img.load()
|
|
parameters = img.text["parameters"]
|
|
|
|
assert parameters == (
|
|
"the prompt of this image\nNegative prompt: bad\n"
|
|
"Steps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51, Size: 16x16"
|
|
)
|
|
# Everything above the settings line is the prompt, exactly what a reader
|
|
# of A1111 parameters text expects.
|
|
assert read_as_a1111_parameters(parameters) == {
|
|
"Steps": "8",
|
|
"Sampler": "euler",
|
|
"Schedule type": "simple",
|
|
"CFG scale": "1",
|
|
"Seed": "51",
|
|
}
|
|
|
|
|
|
def test_path_save_image_with_mask_also_carries_the_sampler_settings(tmp_path):
|
|
"""Both save nodes share the metadata composition."""
|
|
save_node = PathSaveImageRGBA()
|
|
out = str(tmp_path / "a1111_rgba")
|
|
|
|
result = save_node.save_image_with_mask(
|
|
_batch_rgb(1), torch.zeros(1, 16, 16), out, prompt="explicit", workflow_prompt=_sampler_workflow()
|
|
)
|
|
|
|
assert save_success(result) is True
|
|
with Image.open(out + ".png") as img:
|
|
img.load()
|
|
assert img.text["parameters"].endswith("\nSteps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51, Size: 16x16")
|
|
|
|
|
|
def test_path_save_image_sampler_settings_follow_value_links(tmp_path):
|
|
"""A sampler widget is often driven by a plain value node (this pack's
|
|
``IntCreate`` and the like): such a link carries the setting, while a device
|
|
that computes a value stays unknown and is left out rather than guessed."""
|
|
workflow_prompt = {
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"steps": ["5", 0], "cfg": 3.5, "sampler_name": "dpmpp_2m", "seed": ["7", 0], "scheduler": "karras"},
|
|
},
|
|
"5": {"class_type": "Basic data handling: IntCreate", "inputs": {"value": "20"}},
|
|
"7": {"class_type": "RandomInt", "inputs": {"min": 0, "max": 100}},
|
|
}
|
|
out = str(tmp_path / "linked")
|
|
|
|
result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt)
|
|
|
|
assert save_success(result) is True
|
|
with Image.open(out + ".png") as img:
|
|
img.load()
|
|
# "Steps: 20" comes from the value node, the seed cannot be resolved and
|
|
# is therefore missing - which the line itself has to survive.
|
|
assert img.text["parameters"] == "explicit\nSteps: 20, Sampler: dpmpp_2m, Schedule type: karras, CFG scale: 3.5, Size: 16x16"
|
|
assert read_as_a1111_parameters(img.text["parameters"]) == {
|
|
"Steps": "20",
|
|
"Sampler": "dpmpp_2m",
|
|
"Schedule type": "karras",
|
|
"CFG scale": "3.5",
|
|
}
|
|
|
|
|
|
def test_path_save_image_uses_the_last_pass_of_a_multi_sampler_workflow(tmp_path):
|
|
"""A hires fix runs two samplers; like A1111 (which reports its second pass)
|
|
the settings of the pass that produced the file are the ones written."""
|
|
workflow_prompt = {
|
|
"1": _sampler_workflow()["3"],
|
|
"2": {"class_type": "LatentUpscale", "inputs": {"samples": ["1", 0]}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"steps": 20, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "seed": 52, "latent_image": ["2", 0]},
|
|
},
|
|
}
|
|
out = str(tmp_path / "second_pass")
|
|
|
|
result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt)
|
|
|
|
assert save_success(result) is True
|
|
with Image.open(out + ".png") as img:
|
|
img.load()
|
|
assert (
|
|
img.text["parameters"] == "explicit\nSteps: 20, Sampler: dpmpp_2m, Schedule type: karras, CFG scale: 7, Seed: 52, Size: 16x16"
|
|
)
|
|
|
|
|
|
def test_path_save_image_without_sampler_settings_keeps_the_plain_prompt(tmp_path):
|
|
"""A workflow whose settings are spread over several nodes gets no settings
|
|
line at all - the text stays the plain prompt of the previous versions."""
|
|
workflow_prompt = {
|
|
"3": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["4", 0], "guider": ["5", 0], "sampler": ["6", 0]}},
|
|
"4": {"class_type": "BasicScheduler", "inputs": {"steps": 20, "scheduler": "simple"}},
|
|
}
|
|
out = str(tmp_path / "no_settings")
|
|
|
|
result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt)
|
|
|
|
assert save_success(result) is True
|
|
with Image.open(out + ".png") as img:
|
|
img.load()
|
|
assert img.text["parameters"] == "explicit"
|
|
assert read_as_a1111_parameters(img.text["parameters"]) is None
|
|
|
|
|
|
def test_build_generation_parameters_reads_only_what_the_graph_provides():
|
|
"""The settings line is derived from the workflow and never invented."""
|
|
import src.basic_data_handling.path_nodes as pn
|
|
|
|
assert pn.build_generation_parameters(None) == ""
|
|
assert pn.build_generation_parameters("not a workflow") == ""
|
|
assert pn.build_generation_parameters({}) == ""
|
|
# The image size is part of the line only when it is known ...
|
|
assert pn.build_generation_parameters(_sampler_workflow()) == "Steps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51"
|
|
# ... a sampling that only partly denoises is reported like A1111 does ...
|
|
assert pn.build_generation_parameters(_sampler_workflow(denoise=0.6)).endswith("Denoising strength: 0.6")
|
|
# ... and without the settings the line is built around there is no line.
|
|
for missing in ("steps", "cfg", "sampler_name"):
|
|
workflow_prompt = _sampler_workflow()
|
|
del workflow_prompt["3"]["inputs"][missing]
|
|
assert pn.build_generation_parameters(workflow_prompt) == ""
|
|
|
|
|
|
def test_path_save_image_workflow_metadata_in_exif(tmp_path, monkeypatch):
|
|
"""Formats without a native text chunk keep the explicit prompt in
|
|
UserComment and store the workflow in the fields ComfyUI uses for them."""
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: str(tmp_path))
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
workflow_prompt = {"3": {"class_type": "KSampler", "inputs": {}}}
|
|
extra_pnginfo = {"workflow": {"nodes": [{"id": 3}]}}
|
|
out = str(tmp_path / "exif_meta")
|
|
|
|
result = save_node.save_image(
|
|
img, out, format="jpg", prompt="explicit", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo
|
|
)
|
|
assert save_success(result) is True
|
|
|
|
with Image.open(out + ".jpg") as loaded:
|
|
loaded.load()
|
|
exif = loaded.getexif()
|
|
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == "explicit"
|
|
# ComfyUI's own convention for these formats.
|
|
assert exif.get(0x0110) == "prompt:" + json.dumps(workflow_prompt)
|
|
assert exif.get(0x010F) == "workflow:" + json.dumps(extra_pnginfo["workflow"])
|
|
|
|
|
|
def _oversized_workflow_metadata():
|
|
"""A workflow whose full graph JSON comfortably exceeds one JPEG EXIF block
|
|
(a JPEG carries EXIF in a single 16-bit-sized APP1 marker)."""
|
|
big_graph = {"nodes": [{"id": i, "type": "SomeNode", "widgets_values": ["x" * 200]} for i in range(300)]}
|
|
workflow_prompt = {"3": {"class_type": "KSampler", "inputs": {"text": "a" * 500}}}
|
|
return workflow_prompt, {"workflow": big_graph}
|
|
|
|
|
|
def test_path_save_image_oversized_jpeg_metadata_keeps_prompt(tmp_path, capsys):
|
|
"""Too much metadata for a JPEG EXIF block must not fail the save: the
|
|
largest (full-graph) channel is dropped, the rest is kept, and it is
|
|
reported."""
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
workflow_prompt, extra_pnginfo = _oversized_workflow_metadata()
|
|
out = str(tmp_path / "oversized")
|
|
|
|
result = save_node.save_image(
|
|
img, out, format="jpg", prompt="explicit", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo
|
|
)
|
|
assert save_success(result) is True
|
|
assert os.path.exists(out + ".jpg")
|
|
|
|
with Image.open(out + ".jpg") as loaded:
|
|
loaded.load()
|
|
assert loaded.format == "JPEG"
|
|
exif = loaded.getexif()
|
|
# The explicit prompt (and the smaller API prompt) survive ...
|
|
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == "explicit"
|
|
assert exif.get(0x0110) == "prompt:" + json.dumps(workflow_prompt)
|
|
# ... while the oversized full workflow graph was dropped.
|
|
assert 0x010F not in exif
|
|
|
|
assert 'had to drop the full "workflow" graph' in capsys.readouterr().out
|
|
|
|
|
|
def test_path_save_image_oversized_jpeg_drops_all_metadata_but_saves(tmp_path, capsys):
|
|
"""When even the remaining channels do not fit, the image is still written -
|
|
just without metadata."""
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
_, extra_pnginfo = _oversized_workflow_metadata()
|
|
out = str(tmp_path / "huge")
|
|
|
|
# The explicit prompt itself is far beyond a single EXIF block.
|
|
result = save_node.save_image(img, out, format="jpg", prompt="x" * 70000, workflow_extra_pnginfo=extra_pnginfo)
|
|
assert save_success(result) is True
|
|
assert os.path.exists(out + ".jpg")
|
|
|
|
with Image.open(out + ".jpg") as loaded:
|
|
loaded.load()
|
|
assert loaded.format == "JPEG"
|
|
assert not loaded.getexif()
|
|
|
|
assert "saving the image without metadata" in capsys.readouterr().out
|
|
|
|
|
|
def test_path_save_image_png_keeps_oversized_workflow(tmp_path):
|
|
"""Formats with room enough (PNG text chunks) keep the metadata untouched."""
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
workflow_prompt, extra_pnginfo = _oversized_workflow_metadata()
|
|
out = str(tmp_path / "big_png")
|
|
|
|
result = save_node.save_image(img, out, format="png", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo)
|
|
assert save_success(result) is True
|
|
|
|
with Image.open(out + ".png") as loaded:
|
|
loaded.load()
|
|
assert json.loads(loaded.text["workflow"]) == extra_pnginfo["workflow"]
|
|
|
|
|
|
def test_write_pil_image_retries_without_metadata(tmp_path, monkeypatch):
|
|
"""If embedding metadata fails for an unexpected reason, the image is saved
|
|
again without it instead of losing the image."""
|
|
import src.basic_data_handling.path_nodes as pn
|
|
|
|
attempts = []
|
|
|
|
def fake_save(pil_img, path, fmt, quality, metadata_text, workflow_metadata=None):
|
|
attempts.append(bool(metadata_text or workflow_metadata))
|
|
if metadata_text or workflow_metadata:
|
|
raise ValueError("EXIF data is too long")
|
|
pil_img.save(path, format="PNG")
|
|
|
|
monkeypatch.setattr(pn, "_save_pil_image", fake_save)
|
|
img = Image.new("RGB", (8, 8), (255, 0, 0))
|
|
out = str(tmp_path / "retry.png")
|
|
|
|
assert pn._write_pil_image(img, out, "png", 95, "prompt", {"workflow": "{}"}) is True
|
|
assert attempts == [True, False]
|
|
with Image.open(out) as loaded:
|
|
loaded.load()
|
|
|
|
|
|
def test_path_save_image_metadata_can_be_disabled(tmp_path, monkeypatch):
|
|
"""ComfyUI's ``--disable-metadata`` suppresses both metadata channels, like
|
|
it does in the built-in save node."""
|
|
import src.basic_data_handling.path_nodes as pn
|
|
|
|
monkeypatch.setattr(pn, "_metadata_disabled", lambda: True)
|
|
save_node = PathSaveImageRGB()
|
|
img = _batch_rgb(1)
|
|
out = str(tmp_path / "disabled")
|
|
|
|
result = save_node.save_image(img, out, prompt="explicit", workflow_prompt={"1": {}}, workflow_extra_pnginfo={"workflow": {}})
|
|
|
|
assert save_success(result) is True
|
|
with Image.open(out + ".png") as loaded:
|
|
loaded.load()
|
|
assert loaded.text.get("parameters") is None
|
|
assert loaded.text.get("prompt") is None
|
|
assert loaded.text.get("workflow") is None
|
|
|
|
|
|
def test_path_normalize():
|
|
node = PathNormalize()
|
|
assert node.normalize_path("folder/../file.txt") == (os.path.normpath("folder/../file.txt"),)
|
|
assert node.normalize_path("folder//file.txt") == (os.path.normpath("folder//file.txt"),)
|
|
# Test with empty path
|
|
assert node.normalize_path("") == (".",)
|
|
# Test with current directory references
|
|
assert node.normalize_path("./folder/./file.txt") == (os.path.normpath("./folder/./file.txt"),)
|
|
# Test with parent directory references that cancel out
|
|
assert node.normalize_path("folder/../folder/file.txt") == (os.path.normpath("folder/../folder/file.txt"),)
|
|
# Test with mixed separators
|
|
mixed_sep_path = "folder" + os.sep + "subfolder" + ("/" if os.sep == "\\" else "\\") + "file.txt"
|
|
assert node.normalize_path(mixed_sep_path) == (os.path.normpath(mixed_sep_path),)
|
|
|
|
|
|
def test_path_relative(tmp_path):
|
|
node = PathRelative()
|
|
start = tmp_path / "start"
|
|
start.mkdir()
|
|
target = tmp_path / "start/target"
|
|
target.mkdir()
|
|
|
|
assert node.get_relative_path(str(target), str(start)) == (os.path.relpath(str(target), str(start)),)
|
|
assert node.get_relative_path(str(target)) == (os.path.relpath(str(target), os.getcwd()),)
|
|
# Test with same path
|
|
assert node.get_relative_path(str(start), str(start)) == (".",)
|
|
# Test with sibling paths
|
|
sibling = tmp_path / "sibling"
|
|
sibling.mkdir()
|
|
expected = os.path.relpath(str(sibling), str(start))
|
|
assert node.get_relative_path(str(sibling), str(start)) == (expected,)
|
|
# Test with parent directory
|
|
parent_rel = os.path.relpath(str(tmp_path), str(target))
|
|
assert node.get_relative_path(str(tmp_path), str(target)) == (parent_rel,)
|
|
|
|
|
|
def test_path_glob(tmp_path):
|
|
node = PathGlob()
|
|
(tmp_path / "file1.txt").write_text("content")
|
|
(tmp_path / "file2.txt").write_text("content")
|
|
(tmp_path / "file3.txt").write_text("content")
|
|
(tmp_path / "image1.jpg").write_text("image")
|
|
(tmp_path / "image2.png").write_text("image")
|
|
# Create subdirectory with files
|
|
subdir = tmp_path / "subdir"
|
|
subdir.mkdir()
|
|
(subdir / "subfile1.txt").write_text("sub content")
|
|
(subdir / "subfile2.txt").write_text("sub content")
|
|
|
|
# Test simple glob pattern
|
|
result = node.glob_paths(str(tmp_path / "*.txt"))
|
|
expected = sorted([str(tmp_path / f"file{i}.txt") for i in range(1, 4)])
|
|
assert sorted(result[0]) == expected
|
|
|
|
# Test with image files
|
|
img_result = node.glob_paths(str(tmp_path / "image*"))
|
|
img_expected = sorted([str(tmp_path / "image1.jpg"), str(tmp_path / "image2.png")])
|
|
assert sorted(img_result[0]) == img_expected
|
|
|
|
# Test with character class
|
|
char_result = node.glob_paths(str(tmp_path / "file[1-2].txt"))
|
|
char_expected = sorted([str(tmp_path / f"file{i}.txt") for i in range(1, 3)])
|
|
assert sorted(char_result[0]) == char_expected
|
|
|
|
# Test recursive globbing
|
|
recursive_result = node.glob_paths(str(tmp_path / "**/*.txt"), recursive=True)
|
|
recursive_expected = sorted(
|
|
[str(tmp_path / f"file{i}.txt") for i in range(1, 4)] + [str(subdir / f"subfile{i}.txt") for i in range(1, 3)]
|
|
)
|
|
assert sorted(recursive_result[0]) == recursive_expected
|
|
|
|
# Test with no matches
|
|
no_match_result = node.glob_paths(str(tmp_path / "nomatch*.txt"))
|
|
assert len(no_match_result[0]) == 0
|
|
|
|
|
|
def test_path_input_output_dir(monkeypatch):
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_input_directory", lambda: "/tmp/input")
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: "/tmp/output")
|
|
|
|
input_node = PathInputDir()
|
|
output_node = PathOutputDir()
|
|
|
|
assert input_node.execute() == ("/tmp/input",)
|
|
assert output_node.execute() == ("/tmp/output",)
|
|
|
|
# fallback to defaults if functions are not patched
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_input_directory", lambda: "./")
|
|
monkeypatch.setattr("src.basic_data_handling.path_nodes.get_output_directory", lambda: "./")
|
|
|
|
assert PathInputDir().execute() == ("./",)
|
|
assert PathOutputDir().execute() == ("./",)
|
|
|
|
|
|
def test_path_expand_vars(monkeypatch):
|
|
node = PathExpandVars()
|
|
monkeypatch.setenv("TEST_VAR", "test_value")
|
|
monkeypatch.setenv("PATH_VAR", "path/to/somewhere")
|
|
monkeypatch.setenv("EMPTY_VAR", "")
|
|
|
|
assert node.expand_vars("$TEST_VAR/path") == ("test_value/path",)
|
|
assert node.expand_vars("${TEST_VAR}/path") == ("test_value/path",)
|
|
# Test with multiple variables
|
|
assert node.expand_vars("$TEST_VAR/$PATH_VAR") == ("test_value/path/to/somewhere",)
|
|
# Test with empty variable
|
|
assert node.expand_vars("$EMPTY_VAR/path") == ("/path",)
|
|
# Test with no variables
|
|
assert node.expand_vars("regular/path") == ("regular/path",)
|
|
# Test with empty string
|
|
assert node.expand_vars("") == ("",)
|
|
|
|
|
|
def test_path_get_cwd():
|
|
node = PathGetCwd()
|
|
assert node.get_cwd() == (os.getcwd(),)
|
|
# Verify that the result is an absolute path
|
|
assert os.path.isabs(node.get_cwd()[0])
|
|
# Verify that the directory exists
|
|
assert os.path.exists(node.get_cwd()[0])
|
|
assert os.path.isdir(node.get_cwd()[0])
|
|
|
|
|
|
def test_path_list_dir(tmp_path):
|
|
node = PathListDir()
|
|
# Set up a directory structure for testing
|
|
(tmp_path / "file1.txt").write_text("content")
|
|
(tmp_path / "file2.txt").write_text("content")
|
|
(tmp_path / "image.jpg").write_text("image")
|
|
(tmp_path / "dir1").mkdir()
|
|
(tmp_path / "dir2").mkdir()
|
|
(tmp_path / ".hidden_file").write_text("hidden")
|
|
(tmp_path / ".hidden_dir").mkdir()
|
|
|
|
# Test listing all entries
|
|
all_entries = node.list_directory(str(tmp_path))
|
|
assert "file1.txt" in all_entries[0]
|
|
assert "file2.txt" in all_entries[0]
|
|
assert "image.jpg" in all_entries[0]
|
|
assert "dir1" in all_entries[0]
|
|
assert "dir2" in all_entries[0]
|
|
assert ".hidden_file" in all_entries[0]
|
|
assert ".hidden_dir" in all_entries[0]
|
|
|
|
# Test listing only files
|
|
only_files = node.list_directory(str(tmp_path), files_only=True)
|
|
assert "file1.txt" in only_files[0]
|
|
assert "file2.txt" in only_files[0]
|
|
assert "image.jpg" in only_files[0]
|
|
assert ".hidden_file" in only_files[0]
|
|
assert "dir1" not in only_files[0]
|
|
assert "dir2" not in only_files[0]
|
|
assert ".hidden_dir" not in only_files[0]
|
|
|
|
# Test listing only directories
|
|
only_dirs = node.list_directory(str(tmp_path), dirs_only=True)
|
|
assert "file1.txt" not in only_dirs[0]
|
|
assert "file2.txt" not in only_dirs[0]
|
|
assert "image.jpg" not in only_dirs[0]
|
|
assert ".hidden_file" not in only_dirs[0]
|
|
assert "dir1" in only_dirs[0]
|
|
assert "dir2" in only_dirs[0]
|
|
assert ".hidden_dir" in only_dirs[0]
|
|
|
|
# Test with non-existent directory
|
|
with pytest.raises(FileNotFoundError):
|
|
node.list_directory(str(tmp_path / "nonexistent"))
|
|
|
|
# Test with file instead of directory
|
|
with pytest.raises(NotADirectoryError):
|
|
node.list_directory(str(tmp_path / "file1.txt"))
|
|
|
|
# Test with empty path (should use current directory)
|
|
assert len(node.list_directory("")[0]) > 0
|
|
|
|
|
|
def test_path_is_absolute():
|
|
node = PathIsAbsolute()
|
|
# Test with absolute paths
|
|
assert node.check_is_absolute("/") == (True,)
|
|
assert node.check_is_absolute(os.path.abspath(".")) == (True,)
|
|
# Test with relative paths
|
|
assert node.check_is_absolute("relative/path") == (False,)
|
|
assert node.check_is_absolute("./path") == (False,)
|
|
assert node.check_is_absolute("../path") == (False,)
|
|
# Test with empty path
|
|
assert node.check_is_absolute("") == (False,)
|
|
# Test Windows-specific paths if applicable
|
|
if platform.system() == "Windows":
|
|
assert node.check_is_absolute("C:\\Windows") == (True,)
|
|
assert node.check_is_absolute("C:/Windows") == (True,)
|
|
assert node.check_is_absolute("\\\\server\\share") == (True,)
|
|
assert node.check_is_absolute("C:relative\\path") == (False,)
|
|
|
|
|
|
def test_path_common_prefix():
|
|
node = PathCommonPrefix()
|
|
# Test with paths having common prefix
|
|
path1 = "/path/to/file1.txt"
|
|
path2 = "/path/to/file2.txt"
|
|
assert node.get_common_prefix(path1, path2) == ("/path/to/file",)
|
|
|
|
# Test with paths without common prefix
|
|
path3 = "/path/to/file.txt"
|
|
path4 = "/different/path/file.txt"
|
|
assert node.get_common_prefix(path3, path4) == ("/",)
|
|
|
|
# Test with exact same paths
|
|
assert node.get_common_prefix(path1, path1) == (path1,)
|
|
|
|
# Test with one path being a prefix of the other
|
|
path5 = "/path/to"
|
|
path6 = "/path/to/subdir/file.txt"
|
|
assert node.get_common_prefix(path5, path6) == ("/path/to",)
|
|
|
|
# Test with both empty paths
|
|
assert node.get_common_prefix("", "") == ("",)
|