From f931308026d6a9e55d1f0e7ef45c7d876eeaf125 Mon Sep 17 00:00:00 2001 From: David Feltell Date: Mon, 1 Sep 2025 19:32:41 +0100 Subject: [PATCH] [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 --- src/openassetio_comfyui/nodes.py | 176 ++++++++++++++++++++++++++++-- tests/resources/bal_db.json | 38 +++++++ tests/test_openassetio_comfyui.py | 133 +++++++++++++++++++++- 3 files changed, 334 insertions(+), 13 deletions(-) diff --git a/src/openassetio_comfyui/nodes.py b/src/openassetio_comfyui/nodes.py index 9ab92f7..7f759b5 100644 --- a/src/openassetio_comfyui/nodes.py +++ b/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", } diff --git a/tests/resources/bal_db.json b/tests/resources/bal_db.json index eaa6fb8..d443d3b 100644 --- a/tests/resources/bal_db.json +++ b/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" + } + } + } + } } } } diff --git a/tests/test_openassetio_comfyui.py b/tests/test_openassetio_comfyui.py index 7cf0bf6..6eb8b4e 100644 --- a/tests/test_openassetio_comfyui.py +++ b/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): """