[Core] Add PublishImage node

Copy-paste current implementation of core SaveImage, but use a string
input providing an entity reference instead of a file name, and
retrieve target file path from the OpenAssetIO manager.

Support image preview in the UI by copying the published file to a
(temp) directory on ComfyUI's allow-list.

Signed-off-by: David Feltell <david.feltell@foundry.com>

src/openassetio_comfyui/nodes.py

@@ -9,26 +9,35 @@ application.
 """
 
 import hashlib
+import json
 import logging
+import pathlib
+import shutil
 
 import torch
 import numpy as np
 
 from PIL import Image, ImageOps, ImageSequence
+from PIL.PngImagePlugin import PngInfo
 
 from openassetio import EntityReference
 from openassetio.utils import FileUrlPathConverter
 from openassetio.log import LoggerInterface
-from openassetio.access import ResolveAccess
+from openassetio.access import ResolveAccess, PublishingAccess
 from openassetio.hostApi import ManagerFactory, HostInterface
 from openassetio.pluginSystem import (
     HybridPluginSystemManagerImplementationFactory,
     PythonPluginSystemManagerImplementationFactory,
     CppPluginSystemManagerImplementationFactory,
 )
+from openassetio_mediacreation.specifications.twoDimensional import (
+    PlanarBitmapImageResourceSpecification,
+)
 from openassetio_mediacreation.traits.content import LocatableContentTrait
 
+import folder_paths
 import node_helpers
+from comfy.cli_args import args
 
 
 class _OpenAssetIOHost:
@@ -87,19 +96,19 @@ class _OpenAssetIOHost:
         """
         Load and initialise an OpenAssetIO manager plugin.
         """
-        self.__logger = self._LoggerInterface()
+        self.logger = self._LoggerInterface()
         # Initialise plugin system, then find and load a manager.
         self.manager = ManagerFactory.defaultManagerForInterface(
             self._HostInterface(),
             HybridPluginSystemManagerImplementationFactory(
                 # Prefer C++ over Python plugins/methods.
                 [
-                    CppPluginSystemManagerImplementationFactory(self.__logger),
-                    PythonPluginSystemManagerImplementationFactory(self.__logger),
+                    CppPluginSystemManagerImplementationFactory(self.logger),
+                    PythonPluginSystemManagerImplementationFactory(self.logger),
                 ],
-                self.__logger,
+                self.logger,
             ),
-            self.__logger,
+            self.logger,
         )
         if self.manager is None:
             raise RuntimeError(
@@ -107,10 +116,14 @@ class _OpenAssetIOHost:
                 " configuration is correct and that the environment variable"
                 " OPENASSETIO_DEFAULT_CONFIG is set to a valid configuration file."
             )
-        self.__context = self.manager.createContext()
-        self.__file_url_path_converter = FileUrlPathConverter()
+        self.context = self.manager.createContext()
+        self.file_url_path_converter = FileUrlPathConverter()
 
-    def resolve_to_path(self, entity_reference: str) -> str:
+    def resolve_to_path(
+        self,
+        entity_reference: EntityReference | str,
+        access_mode: ResolveAccess = ResolveAccess.kRead,
+    ) -> str:
         """
         Resolve an OpenAssetIO entity reference to a local file path.
 
@@ -121,10 +134,11 @@ class _OpenAssetIOHost:
         @throw RuntimeError if the entity reference resolves
         successfully, but the entity has no location.
         """
-        entity_reference = self.manager.createEntityReference(entity_reference)
+        if isinstance(entity_reference, str):
+            entity_reference = self.manager.createEntityReference(entity_reference)
 
         traits_data = self.manager.resolve(
-            entity_reference, {LocatableContentTrait.kId}, ResolveAccess.kRead, self.__context
+            entity_reference, {LocatableContentTrait.kId}, access_mode, self.context
         )
         locatable_content_trait = LocatableContentTrait(traits_data)
         url = locatable_content_trait.getLocation()
@@ -134,7 +148,7 @@ class _OpenAssetIOHost:
                 f"Failed to resolve entity reference '{entity_reference}' to a location"
             )
 
-        return self.__file_url_path_converter.pathFromUrl(url)
+        return self.file_url_path_converter.pathFromUrl(url)
 
 
 class ResolveImage:
@@ -254,12 +268,150 @@ class ResolveImage:
         return output_image, output_mask
 
 
+class PublishImage:
+    """
+    Node to publish an image to an OpenAssetIO entity reference.
+
+    The non-OpenAssetIO logic is largely duplicated from the built-in
+    ComfyUI SaveImage node (at least as of 4f5812b9).
+    """
+
+    # Tooltip to display when hovering over the node.
+    DESCRIPTION = "Publishes images to an asset manager."
+    # Menu category.
+    CATEGORY = "image"
+    # Function to call when node is executed.
+    FUNCTION = "publish_images"
+    # Node outputs.
+    RETURN_TYPES = ()
+    # Marks this node as a terminal node, ensuring the associated
+    # subgraph is executed when running the graph.
+    OUTPUT_NODE = True
+
+    @classmethod
+    def INPUT_TYPES(cls) -> dict:
+        """
+        Input sockets and widgets.
+        """
+        return {
+            "required": {
+                "entity_reference": (
+                    "STRING",
+                    {"multiline": False, "tooltip": "The entity to publish to"},
+                ),
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, entity_reference: str) -> bool:
+        """
+        Validate that the input entity reference is valid syntax for the
+        current manager.
+        """
+        return _OpenAssetIOHost.instance().manager.isEntityReferenceString(entity_reference)
+
+    def __init__(self):
+        """
+        Initialise node.
+
+        Duplicated from SaveImage node.
+        """
+        self.compress_level = 4
+
+    def publish_images(
+        self, entity_reference: str, images: torch.Tensor, prompt=None, extra_pnginfo=None
+    ) -> dict:
+        """
+        Publish the input images to the specified OpenAssetIO entity.
+
+        Assumes the working reference provided by `preflight()` can be
+        resolved to a destination file path to write to.
+
+        Non-OpenAssetIO logic largely duplicated from the built-in
+        ComfyUI SaveImage node (at least as of 4f5812b9).
+
+        In particular, we inherit support for a "%batch_num%"
+        placeholder in the resolved path.
+        """
+        entity_reference = _OpenAssetIOHost.instance().manager.createEntityReference(
+            entity_reference
+        )
+
+        spec = PlanarBitmapImageResourceSpecification.create()
+
+        working_ref = _OpenAssetIOHost.instance().manager.preflight(
+            entity_reference,
+            spec.traitsData(),
+            PublishingAccess.kWrite,
+            _OpenAssetIOHost.instance().context,
+        )
+
+        # Get destination file path (potentially containing batch_num
+        # placeholder). This may be a temporary/staging path, or it
+        # may be the final path, depending on the manager's
+        # implementation.
+        file_path_tmplt = _OpenAssetIOHost.instance().resolve_to_path(
+            working_ref, access_mode=ResolveAccess.kManagerDriven
+        )
+
+        results = list()
+        for batch_number, image in enumerate(images):
+            i = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            metadata = None
+            if not args.disable_metadata:
+                metadata = PngInfo()
+                if prompt is not None:
+                    metadata.add_text("prompt", json.dumps(prompt))
+                if extra_pnginfo is not None:
+                    for x in extra_pnginfo:
+                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+
+            file_path = file_path_tmplt.replace("%batch_num%", str(batch_number))
+
+            img.save(file_path, pnginfo=metadata, compress_level=self.compress_level)
+
+            # Configure traits to register with the asset manager.
+            url = _OpenAssetIOHost.instance().file_url_path_converter.pathToUrl(file_path)
+            spec.locatableContentTrait().setLocation(url)
+
+            # Publish the image to the working reference.
+            final_ref = _OpenAssetIOHost.instance().manager.register(
+                working_ref,
+                spec.traitsData(),
+                PublishingAccess.kWrite,
+                _OpenAssetIOHost.instance().context,
+            )
+            # Get the path of the published image. This may be different
+            # to the path resolved from the working reference (i.e. if
+            # the manager moved it as part of the publishing process).
+            final_file_path = pathlib.Path(_OpenAssetIOHost.instance().resolve_to_path(final_ref))
+            # Copy to ComfyUI temp directory for display in the UI. For
+            # security reasons, ComfyUI does not allow images to be
+            # served from arbitrary paths on disk, so we must copy them
+            # to an allowed location. Here, we choose ComfyUI's temp
+            # directory.
+            shutil.copy2(final_file_path, folder_paths.get_temp_directory())
+
+            results.append({"filename": final_file_path.name, "subfolder": "", "type": "temp"})
+
+        # For output nodes, we can return a dict with a "ui" key,
+        # containing data to display in the ComfyUI interface.
+        # Here, we return the list of published images, which will be
+        # displayed in the node.
+        return {"ui": {"images": results}}
+
+
 # Plugin registration: node classes.
 NODE_CLASS_MAPPINGS = {
     "OpenAssetIOResolveImage": ResolveImage,
+    "OpenAssetIOPublishImage": PublishImage,
 }
 
 # Plugin registration: node names.
 NODE_DISPLAY_NAME_MAPPINGS = {
     "OpenAssetIOResolveImage": "OpenAssetIO Resolve Image",
+    "OpenAssetIOPublishImage": "OpenAssetIO Publish Image",
 }

tests/resources/bal_db.json

@@ -1,4 +1,21 @@
 {
+  "managementPolicy": {
+    "read": {
+      "default": {
+        "openassetio-mediacreation:managementPolicy.Managed": {}
+      }
+    },
+    "write": {
+      "default": {
+        "openassetio-mediacreation:managementPolicy.Managed": {}
+      }
+    },
+    "managerDriven": {
+      "default": {
+        "openassetio-mediacreation:managementPolicy.Managed": {}
+      }
+    }
+  },
   "entities": {
     "image_file": {
       "versions": [
@@ -22,6 +39,27 @@
           }
         }
       ]
+    },
+    "new_image_file": {
+      "versions": [],
+      "overrideByAccess": {
+        "write": {
+          "traits": {
+            "openassetio-mediacreation:usage.Entity": {},
+            "openassetio-mediacreation:twoDimensional.Image": {},
+            "openassetio-mediacreation:twoDimensional.PixelBased": {},
+            "openassetio-mediacreation:twoDimensional.Planar": {},
+            "openassetio-mediacreation:content.LocatableContent": {}
+          }
+        },
+        "managerDriven": {
+          "traits": {
+            "openassetio-mediacreation:content.LocatableContent": {
+              "location": "file://${test_tmp_dir}/dog.png"
+            }
+          }
+        }
+      }
     }
   }
 }

tests/test_openassetio_comfyui.py

@@ -19,7 +19,9 @@ import pytest
 import torch
 from PIL import Image
 
-from openassetio_comfyui.nodes import ResolveImage, _OpenAssetIOHost
+from openassetio_comfyui.nodes import ResolveImage, _OpenAssetIOHost, PublishImage
+
+import folder_paths
 
 
 class Test_ResolveImage_init:
@@ -131,6 +133,102 @@ class Test_ResolveImage_resolve_image:
             resolve_image_node.resolve_image("bal:///image_file")
 
 
+class Test_PublishImage_init:
+    """
+    Test that the node can be instantiated.
+    """
+
+    def test_has_correct_type(self, publish_image_node):
+        assert isinstance(publish_image_node, PublishImage)
+
+
+class Test_PublishImage_constants:
+    """
+    Test the node's metadata.
+    """
+
+    def test_has_expected_values(self):
+        assert PublishImage.DESCRIPTION == "Publishes images to an asset manager."
+        assert PublishImage.CATEGORY == "image"
+        assert PublishImage.FUNCTION == "publish_images"
+        assert PublishImage.RETURN_TYPES == tuple()
+        assert PublishImage.OUTPUT_NODE is True
+
+    def test_function_matches(self):
+        assert inspect.isfunction(getattr(PublishImage, PublishImage.FUNCTION))
+
+
+class Test_PublishImage_INPUT_TYPES:
+    def test_has_expected_structure(self):
+        input_types = PublishImage.INPUT_TYPES()
+        assert input_types == {
+            "required": {
+                "entity_reference": (
+                    "STRING",
+                    {"multiline": False, "tooltip": "The entity to publish to"},
+                ),
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+
+class Test_PublishImage_VALIDATE_INPUTS:
+    def test_when_is_a_reference_then_returns_true(self):
+        valid = PublishImage.VALIDATE_INPUTS("bal:///")
+        assert valid
+
+    def test_when_not_a_reference_then_returns_false(self):
+        valid = PublishImage.VALIDATE_INPUTS("notbal:///")
+        assert not valid
+
+    def test_when_no_manager_found_then_raises(self, assert_raises_missing_manager):
+        with assert_raises_missing_manager():
+            PublishImage.VALIDATE_INPUTS("bal:///")
+
+
+class Test_PublishImage_publish_images:
+    def test_when_valid_reference_then_writes_image(
+        self, publish_image_node, image_file_path, image_file_tensor
+    ):
+        publish_image_node.publish_images("bal:///new_image_file", image_file_tensor)
+
+        new_image_file_path = image_file_path.with_name("dog.png")  # See bal_db.json
+        assert new_image_file_path.exists()
+
+    def test_when_valid_reference_then_provides_preview_of_result(
+        self, publish_image_node, image_file_tensor, tmp_path
+    ):
+        expected_ui_desc = {
+            "ui": {
+                "images": [
+                    {
+                        "type": "temp",
+                        "subfolder": "",
+                        "filename": "dog.png",  # See bal_db.json
+                    }
+                ]
+            }
+        }
+        actual_ui_desc = publish_image_node.publish_images(
+            "bal:///new_image_file", image_file_tensor
+        )
+
+        assert expected_ui_desc == actual_ui_desc
+
+        published_file_path = tmp_path / "dog.png"
+        ui_file_path = pathlib.Path(folder_paths.get_temp_directory()) / "dog.png"
+
+        assert ui_file_path != published_file_path  # Confidence check.
+        assert ui_file_path.read_bytes() == published_file_path.read_bytes()  # i.e. copied.
+
+    def test_when_no_manager_found_then_raises(
+        self, publish_image_node, assert_raises_missing_manager
+    ):
+        with assert_raises_missing_manager():
+            publish_image_node.publish_images("bal:///new_image_file", torch.zeros((1, 2, 2, 3)))
+
+
 @pytest.fixture
 def assert_raises_missing_manager(monkeypatch):
     """
@@ -156,6 +254,17 @@ def assert_raises_missing_manager(monkeypatch):
     return assert_ctx
 
 
+@pytest.fixture
+def image_file_tensor(resolve_image_node, image_file_path):
+    """
+    Fixture to provide a tensor representation of the input image file.
+    """
+    # Short-cut making use of ResolveImage node to get a tensor with
+    # expected structure.
+    images, _ = resolve_image_node.resolve_image("bal:///image_file")
+    return images
+
+
 @pytest.fixture
 def image_file_hash(image_file_path):
     """
@@ -184,6 +293,14 @@ def image_file_path(monkeypatch, tmp_path):
     return file_path
 
 
+@pytest.fixture
+def publish_image_node():
+    """
+    Fixture to create a PublishImage node instance.
+    """
+    return PublishImage()
+
+
 @pytest.fixture
 def resolve_image_node():
     """
@@ -192,6 +309,20 @@ def resolve_image_node():
     return ResolveImage()
 
 
+@pytest.fixture(autouse=True)
+def set_comfyui_temp_dir(tmp_path_factory):
+    """
+    Override default ComfyUI temp directory to a location we control.
+
+    We do this because:
+    * ComfyUI's default temp directory (a top-level "temp" directory
+      under the ComfyUI repo) doesn't exist during test runs.
+    * We would like to ensure cleanup when the test completes.
+    """
+    tmp_dir = tmp_path_factory.mktemp("temp")
+    folder_paths.set_temp_directory(str(tmp_dir))
+
+
 @pytest.fixture(autouse=True)
 def openassetio_config_env_var(monkeypatch, resources_dir):
     """