nodes/StableLlama-ComfyUI-basic_data_handling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
StableLlama-ComfyUI-basic_d…/tests/test_path_nodes.py
T

1478 lines
59 KiB
Python
Raw Permalink Blame History

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 &lt;lora:x:1.0&gt;\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("", "") == ("",)
Reference in New Issue View Git Blame Copy Permalink