[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 <[email protected]>

Author: feltech

src/openassetio_comfyui/nodes.py

@@ -9,26 +9,35 @@
 """
 
 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,30 +96,34 @@ def __init__(self):
         """
         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(
                 "Could not create an OpenAssetIO manager instance. Ensure that your OpenAssetIO"
                 " 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()
-
-    def resolve_to_path(self, entity_reference: str) -> str:
+        self.context = self.manager.createContext()
+        self.file_url_path_converter = FileUrlPathConverter()
+
+    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 @@ def resolve_to_path(self, entity_reference: str) -> str:
         @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 @@ def resolve_to_path(self, entity_reference: str) -> str:
                 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 @@ def resolve_image(self, entity_reference: str) -> tuple[torch.Tensor, torch.Tens
         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"
+            }
+          }
+        }
+      }
     }
   }
 }