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 " 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 \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"" 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"] 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 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("", "") == ("",)