Add nodes to support OpenAPI and similar backend workflows

tests/compare/conftest.py

@@ -29,6 +29,8 @@ def args_pytest(pytestconfig):
 
 def gather_file_basenames(directory: str):
     files = []
+    if not os.path.isdir(directory):
+        return files
     for file in os.listdir(directory):
         if file.endswith(".png"):
             files.append(file)

tests/compare/test_quality.py

@@ -1,6 +1,8 @@
 import datetime
 import numpy as np
 import os
+
+import torch
 from PIL import Image
 import pytest
 from pytest import fixture
@@ -9,12 +11,13 @@ from typing import Tuple, List
 from cv2 import imread, cvtColor, COLOR_BGR2RGB
 from skimage.metrics import structural_similarity as ssim
 
-
 """
 This test suite compares images in 2 directories by file name
 The directories are specified by the command line arguments --baseline_dir and --test_dir
 
 """
+
+
 # ssim: Structural Similarity Index
 # Returns a tuple of (ssim, diff_image)
 def ssim_score(img0: np.ndarray, img1: np.ndarray) -> Tuple[float, np.ndarray]:
@@ -22,7 +25,8 @@ def ssim_score(img0: np.ndarray, img1: np.ndarray) -> Tuple[float, np.ndarray]:
     # rescale the difference image to 0-255 range
     diff = (diff * 255).astype("uint8")
     return score, diff
-    
+
+
 # Metrics must return a tuple of (score, diff_image)
 METRICS = {"ssim": ssim_score}
 METRICS_PASS_THRESHOLD = {"ssim": 0.95}
@@ -32,7 +36,7 @@ class TestCompareImageMetrics:
     @fixture(scope="class")
     def test_file_names(self, args_pytest):
         test_dir = args_pytest['test_dir']
-        fnames = self.gather_file_basenames(test_dir)  
+        fnames = self.gather_file_basenames(test_dir)
         yield fnames
         del fnames
 
@@ -56,50 +60,53 @@ class TestCompareImageMetrics:
                     score = self.lookup_score_from_fname(file, metrics_file)
                     image_file_list = []
                     image_file_list.append([
-                                            os.path.join(baseline_dir, file),
-                                            os.path.join(test_dir, file),
-                                            os.path.join(metric_path, file)
-                                            ])
+                        os.path.join(baseline_dir, file),
+                        os.path.join(test_dir, file),
+                        os.path.join(metric_path, file)
+                    ])
                     # Create grid
                     image_list = [[Image.open(file) for file in files] for files in image_file_list]
                     grid = self.image_grid(image_list)
                     grid.save(os.path.join(grid_dir, f"{metric_dir}_{score:.3f}_{file}"))
-    
+
     # Tests run for each baseline file name
     @fixture()
     def fname(self, baseline_fname, teardown):
         yield baseline_fname
         del baseline_fname
-    
-    def test_directories_not_empty(self, args_pytest, teardown):
-        baseline_dir = args_pytest['baseline_dir']
-        test_dir = args_pytest['test_dir']
-        assert len(os.listdir(baseline_dir)) != 0, f"Baseline directory {baseline_dir} is empty"
-        assert len(os.listdir(test_dir)) != 0, f"Test directory {test_dir} is empty"
 
-    def test_dir_has_all_matching_metadata(self, fname, test_file_names, args_pytest, teardown):
+    # For a baseline image file, finds the corresponding file name in test_dir and 
+    # compares the images using the metrics in METRICS
+    @pytest.mark.parametrize("metric", METRICS.keys())
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+    def test_pipeline_compare(
+            self,
+            args_pytest,
+            fname,
+            test_file_names,
+            metric,
+            teardown,
+    ):
+        baseline_dir = args_pytest['baseline_dir']
+
+        test_dir = args_pytest['test_dir']
+        metrics_output_file = args_pytest['metrics_file']
+        img_output_dir = args_pytest['img_output_dir']
+
+        if not os.path.isdir(baseline_dir):
+            pytest.skip("Baseline directory does not exist")
+            return
+
+        if not os.path.isdir(test_dir):
+            pytest.skip("Test directory does not exist")
+            return
+
         # Check that all files in baseline_dir have a file in test_dir with matching metadata
         baseline_file_path = os.path.join(args_pytest['baseline_dir'], fname)
         file_paths = [os.path.join(args_pytest['test_dir'], f) for f in test_file_names]
         file_match = self.find_file_match(baseline_file_path, file_paths)
         assert file_match is not None, f"Could not find a file in {args_pytest['test_dir']} with matching metadata to {baseline_file_path}"
 
-    # For a baseline image file, finds the corresponding file name in test_dir and 
-    # compares the images using the metrics in METRICS
-    @pytest.mark.parametrize("metric", METRICS.keys())
-    def test_pipeline_compare(
-        self,
-        args_pytest,
-        fname,
-        test_file_names,
-        metric,
-        teardown,
-    ):
-        baseline_dir = args_pytest['baseline_dir']
-        test_dir = args_pytest['test_dir']
-        metrics_output_file = args_pytest['metrics_file']
-        img_output_dir = args_pytest['img_output_dir']
-        
         baseline_file_path = os.path.join(baseline_dir, fname)
 
         # Find file match
@@ -109,7 +116,7 @@ class TestCompareImageMetrics:
         # Run metrics
         sample_baseline = self.read_img(baseline_file_path)
         sample_secondary = self.read_img(test_file)
-        
+
         score, metric_img = METRICS[metric](sample_baseline, sample_secondary)
         metric_status = score > METRICS_PASS_THRESHOLD[metric]
 
@@ -140,17 +147,17 @@ class TestCompareImageMetrics:
         cols = len(img_list[0])
 
         w, h = img_list[0][0].size
-        grid = Image.new('RGB', size=(cols*w, rows*h))
-        
+        grid = Image.new('RGB', size=(cols * w, rows * h))
+
         for i, row in enumerate(img_list):
             for j, img in enumerate(row):
-                grid.paste(img, box=(j*w, i*h))
+                grid.paste(img, box=(j * w, i * h))
         return grid
 
     def lookup_score_from_fname(self,
                                 fname: str,
                                 metrics_output_file: str
-        ) -> float:
+                                ) -> float:
         fname_basestr = os.path.splitext(fname)[0]
         with open(metrics_output_file, 'r') as f:
             for line in f:
@@ -166,12 +173,12 @@ class TestCompareImageMetrics:
                 files.append(file)
         return files
 
-    def read_file_prompt(self, fname:str) -> str:
+    def read_file_prompt(self, fname: str) -> str:
         # Read prompt from image file metadata
         img = Image.open(fname)
         img.load()
         return img.info['prompt']
-    
+
     def find_file_match(self, baseline_file: str, file_paths: List[str]):
         # Find a file in file_paths with matching metadata to baseline_file
         baseline_prompt = self.read_file_prompt(baseline_file)
@@ -193,4 +200,4 @@ class TestCompareImageMetrics:
         for f in file_paths:
             test_file_prompt = self.read_file_prompt(f)
             if baseline_prompt == test_file_prompt:
-                return f
+                return f

tests/inference/test_inference.py

@@ -148,12 +148,12 @@ scheduler_list = SCHEDULER_NAMES[:]
 @pytest.mark.parametrize("sampler", sampler_list)
 @pytest.mark.parametrize("scheduler", scheduler_list)
 @pytest.mark.parametrize("prompt", prompt_list)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
 class TestInference:
     #
     # Initialize server and client
     #
 
-
     def start_client(self, listen: str, port: int):
         # Start client
         comfy_client = ComfyClient()

tests/nodes/test_openapi_unit.py

@@ -0,0 +1,218 @@
+import os
+import pathlib
+import re
+import uuid
+from datetime import datetime
+
+import pytest
+import torch
+from PIL import Image
+from freezegun import freeze_time
+from comfy.cmd import folder_paths
+from comfy_extras.nodes.nodes_open_api import SaveImagesResponse, IntRequestParameter, FloatRequestParameter, StringRequestParameter, HashImage, StringPosixPathJoin, LegacyOutputURIs, DevNullUris, StringJoin, StringToUri, UriFormat, ImageExifMerge, ImageExifCreationDateAndBatchNumber, ImageExif, ImageExifUncommon
+
+_image_1x1 = torch.zeros((1, 1, 3), dtype=torch.float32, device="cpu")
+
+
+@pytest.fixture(scope="function", autouse=True)
+def use_tmp_path(tmp_path: pathlib.Path):
+    orig_dir = folder_paths.get_output_directory()
+    folder_paths.set_output_directory(tmp_path)
+    yield tmp_path
+    folder_paths.set_output_directory(orig_dir)
+
+
+def test_save_image_response():
+    assert SaveImagesResponse.INPUT_TYPES() is not None
+    n = SaveImagesResponse()
+    result = n.execute(images=[_image_1x1], uris=["with_prefix/1.png"], name="test")
+    assert os.path.isfile(os.path.join(folder_paths.get_output_directory(), "with_prefix/1.png"))
+    assert len(result["result"]) == 1
+    assert len(result["ui"]["images"]) == 1
+    assert result["result"][0]["filename"] == "1.png"
+    assert result["result"][0]["subfolder"] == "with_prefix"
+    assert result["result"][0]["name"] == "test"
+
+
+def test_save_image_response_abs_local_uris():
+    assert SaveImagesResponse.INPUT_TYPES() is not None
+    n = SaveImagesResponse()
+    result = n.execute(images=[_image_1x1], uris=[os.path.join(folder_paths.get_output_directory(), "with_prefix/1.png")], name="test")
+    assert os.path.isfile(os.path.join(folder_paths.get_output_directory(), "with_prefix/1.png"))
+    assert len(result["result"]) == 1
+    assert len(result["ui"]["images"]) == 1
+    assert result["result"][0]["filename"] == "1.png"
+    assert result["result"][0]["subfolder"] == "with_prefix"
+    assert result["result"][0]["name"] == "test"
+
+
+def test_save_image_response_remote_uris():
+    n = SaveImagesResponse()
+    uri = "memory://some_folder/1.png"
+    result = n.execute(images=[_image_1x1], uris=[uri])
+    assert len(result["result"]) == 1
+    assert len(result["ui"]["images"]) == 1
+    filename_ = result["result"][0]["filename"]
+    assert filename_ != "1.png"
+    assert filename_ != ""
+    assert uuid.UUID(filename_.replace(".png", "")) is not None
+    assert os.path.isfile(os.path.join(folder_paths.get_output_directory(), filename_))
+    assert result["result"][0]["abs_path"] == uri
+    assert result["result"][0]["subfolder"] == ""
+
+
+def test_save_exif():
+    n = SaveImagesResponse()
+    filename = "with_prefix/2.png"
+    result = n.execute(images=[_image_1x1], uris=[filename], name="test", exif=[{
+        "Title": "test title"
+    }])
+    filepath = os.path.join(folder_paths.get_output_directory(), filename)
+    assert os.path.isfile(filepath)
+    with Image.open(filepath) as img:
+        assert img.info['Title'] == "test title"
+
+
+def test_no_local_file():
+    n = SaveImagesResponse()
+    uri = "memory://some_folder/2.png"
+    result = n.execute(images=[_image_1x1], uris=[uri], local_uris=["/dev/null"])
+    assert len(result["result"]) == 1
+    assert len(result["ui"]["images"]) == 1
+    assert result["result"][0]["filename"] == ""
+    assert not os.path.isfile(os.path.join(folder_paths.get_output_directory(), result["result"][0]["filename"]))
+    assert result["result"][0]["abs_path"] == uri
+    assert result["result"][0]["subfolder"] == ""
+
+
+def test_int_request_parameter():
+    nt = IntRequestParameter.INPUT_TYPES()
+    assert nt is not None
+    n = IntRequestParameter()
+    v, = n.execute(value=1, name="test")
+    assert v == 1
+
+
+def test_float_request_parameter():
+    nt = FloatRequestParameter.INPUT_TYPES()
+    assert nt is not None
+    n = FloatRequestParameter()
+    v, = n.execute(value=3.5, name="test", description="")
+    assert v == 3.5
+
+
+def test_string_request_parameter():
+    nt = StringRequestParameter.INPUT_TYPES()
+    assert nt is not None
+    n = StringRequestParameter()
+    v, = n.execute(value="test", name="test")
+    assert v == "test"
+
+
+def test_hash_images():
+    nt = HashImage.INPUT_TYPES()
+    assert nt is not None
+    n = HashImage()
+    hashes = n.execute(images=[_image_1x1.clone(), _image_1x1.clone()])
+    # same image, same hash
+    assert hashes[0] == hashes[1]
+    # hash should be a valid sha256 hash
+    p = re.compile(r'^[0-9a-fA-F]{64}$')
+    for hash in hashes:
+        assert p.match(hash)
+
+
+def test_string_posix_path_join():
+    nt = StringPosixPathJoin.INPUT_TYPES()
+    assert nt is not None
+    n = StringPosixPathJoin()
+    joined_path = n.execute(value2="c", value0="a", value1="b")
+    assert joined_path == "a/b/c"
+
+
+def test_legacy_output_uris(use_tmp_path):
+    nt = LegacyOutputURIs.INPUT_TYPES()
+    assert nt is not None
+    n = LegacyOutputURIs()
+    images_ = [_image_1x1, _image_1x1]
+    output_paths = n.execute(images=images_)
+    # from SaveImage node
+    full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("ComfyUI", str(use_tmp_path), images_[0].shape[1], images_[0].shape[0])
+    file1 = f"{filename}_{counter:05}_.png"
+    file2 = f"{filename}_{counter + 1:05}_.png"
+    files = [file1, file2]
+    assert os.path.basename(output_paths[0]) == files[0]
+    assert os.path.basename(output_paths[1]) == files[1]
+
+
+def test_null_uris():
+    nt = DevNullUris.INPUT_TYPES()
+    assert nt is not None
+    n = DevNullUris()
+    res = n.execute([_image_1x1, _image_1x1])
+    assert all(x == "/dev/null" for x in res)
+
+
+def test_string_join():
+    assert StringJoin.INPUT_TYPES() is not None
+    n = StringJoin()
+    assert n.execute(separator="*", value1="b", value3="c", value0="a") == "a*b*c"
+
+
+def test_string_to_uri():
+    assert StringToUri.INPUT_TYPES() is not None
+    n = StringToUri()
+    assert n.execute("x", batch=3) == ["x"] * 3
+
+
+def test_uri_format(use_tmp_path):
+    assert UriFormat.INPUT_TYPES() is not None
+    n = UriFormat()
+    images = [_image_1x1, _image_1x1]
+    # with defaults
+    uris, metadata_uris = n.execute(images=images, uri_template="{output}/{uuid}_{batch_index:05d}.png")
+    for uri in uris:
+        assert os.path.isabs(uri), "uri format returns absolute URIs when output appears"
+        assert os.path.commonpath([uri, use_tmp_path]) == str(use_tmp_path), "should be under output dir"
+    uris, metadata_uris = n.execute(images=images, uri_template="{output}/{uuid}.png")
+    for uri in uris:
+        assert os.path.isabs(uri)
+        assert os.path.commonpath([uri, use_tmp_path]) == str(use_tmp_path), "should be under output dir"
+
+    with pytest.raises(KeyError):
+        n.execute(images=images, uri_template="{xyz}.png")
+
+
+def test_image_exif_merge():
+    assert ImageExifMerge.INPUT_TYPES() is not None
+    n = ImageExifMerge()
+    res = n.execute(value0=[{"a": "1"}, {"a": "1"}], value1=[{"b": "2"}, {"a": "1"}], value2=[{"a": 3}, {}], value4=[{"a": ""}, {}])
+    assert res[0]["a"] == 3
+    assert res[0]["b"] == "2"
+    assert res[1]["a"] == "1"
+
+
+@freeze_time("2012-01-14 03:21:34", tz_offset=-4)
+def test_image_exif_creation_date_and_batch_number():
+    assert ImageExifCreationDateAndBatchNumber.INPUT_TYPES() is not None
+    n = ImageExifCreationDateAndBatchNumber()
+    res = n.execute(images=[_image_1x1, _image_1x1])
+    mock_now = datetime(2012, 1, 13, 23, 21, 34)
+
+    now_formatted = mock_now.strftime("%Y:%m:%d %H:%M:%S%z")
+    assert res[0]["ImageNumber"] == "0"
+    assert res[1]["ImageNumber"] == "1"
+    assert res[0]["CreationDate"] == res[1]["CreationDate"] == now_formatted
+
+
+def test_image_exif():
+    assert ImageExif.INPUT_TYPES() is not None
+    n = ImageExif()
+    res = n.execute(images=[_image_1x1], Title="test", Artist="test2")
+    assert res[0]["Title"] == "test"
+    assert res[0]["Artist"] == "test2"
+
+
+def test_image_exif_uncommon():
+    assert "DigitalZoomRatio" in ImageExifUncommon.INPUT_TYPES()
+    ImageExifUncommon().execute(images=[_image_1x1])