diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..fe4eb9e --- /dev/null +++ b/.editorconfig @@ -0,0 +1,34 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +# http://editorconfig.org + +root = true + +[*] +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true +charset = utf-8 +end_of_line = lf +max_line_length = 99 +# For docstrings and comments. +ij_visual_guides = 72 + +[*.bat] +indent_style = tab +end_of_line = crlf + +[LICENSE] +insert_final_newline = false + +[Makefile] +indent_style = tab + +[{*.yml,*.yaml}] +indent_size = 2 + +[*.json] +indent_size = 2 \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..243396f --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,15 @@ +* openassetio-comfyui version: +* Python version: +* Operating System: + +### Description + +Describe what you were trying to get done. +Tell us what happened, what went wrong, and what you expected to happen. + +### What I Did + +``` +Paste the command(s) you ran and the output. +If there was a crash, please include the traceback here. +``` diff --git a/.github/workflows/build_pipeline.yml b/.github/workflows/build_pipeline.yml new file mode 100644 index 0000000..68adf8f --- /dev/null +++ b/.github/workflows/build_pipeline.yml @@ -0,0 +1,47 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +name: CI build + +on: + pull_request: + workflow_dispatch: + +jobs: + test: + name: Test and Lint + runs-on: ${{ matrix.os }} + env: + PYTHONIOENCODING: "utf8" + strategy: + matrix: + os: [ ubuntu-latest ] + python-version: [ "3.11" ] + + steps: + - uses: actions/checkout@v4 + + - uses: actions/checkout@v4 + name: Checkout ComfyUI + with: + repository: comfyanonymous/ComfyUI + ref: v0.3.57 + path: comfyui + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[dev] + pip install -r comfyui/requirements.txt + - name: Run Tests + run: | + pytest tests/ + env: + PYTHONPATH: ./comfyui + - name: Run Linting + run: | + ruff check . diff --git a/.github/workflows/publish_node.yml b/.github/workflows/publish_node.yml new file mode 100644 index 0000000..8d9d80b --- /dev/null +++ b/.github/workflows/publish_node.yml @@ -0,0 +1,25 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +name: 📦 Publish to Comfy registry +on: + workflow_dispatch: + push: + tags: + - '*' + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: ♻️ Check out code + uses: actions/checkout@v4 + - name: 📦 Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.github/workflows/validate.yml.disabled b/.github/workflows/validate.yml.disabled new file mode 100644 index 0000000..21d89b9 --- /dev/null +++ b/.github/workflows/validate.yml.disabled @@ -0,0 +1,18 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +# TODO(DF): Re-enable once initial PR merged to set a baseline. +name: Validate backwards compatibility + +on: + pull_request: + branches: + - master + - main + +jobs: + validate: + runs-on: ubuntu-latest + steps: + - uses: comfy-org/node-diff@main diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..035b179 --- /dev/null +++ b/.gitignore @@ -0,0 +1,118 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# OSX useful to ignore +*.DS_Store +.AppleDouble +.LSOverride + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log + +# Sphinx documentation +docs/_build/ + +# IntelliJ Idea +.idea +*.iml +*.ipr +*.iws + +# PyBuilder +target/ + +# Cookiecutter +output/ +python_boilerplate/ +cookiecutter-pypackage-env/ + +# vscode settings +.history/ +*.code-workspace + +# Frontend extension +node_modules/ +.env +.env.local +.env.development.local +.env.test.local +.env.production.local +npm-debug.log* +yarn-debug.log* +yarn-error.log* +node.zip +.vscode/ +.claude/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..91a7860 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,14 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.4.9 + hooks: + # Run the linter. + - id: ruff + args: [ --fix ] + # Run the formatter. + - id: ruff-format diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..9f74509 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,13 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +include LICENSE +include README.md + +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif + + diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..cf752c8 --- /dev/null +++ b/NOTICE @@ -0,0 +1,5 @@ +openassetio-comfyui +Copyright (c) 2025 The Foundry Visionmongers Ltd. + +This product includes software developed at +The Foundry Visionmongers Ltd. diff --git a/README.md b/README.md new file mode 100644 index 0000000..abb962f --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +# OpenAssetIO-ComfyUI + +## What + +Custom [ComfyUI](https://comfy.org) nodes for resolving and publishing +assets directly from a workflow via +[OpenAssetIO](https://docs.openassetio.org/OpenAssetIO/). + +## Why + +This project allows ComfyUI to leverage the abilities of +OpenAssetIO-enabled asset management systems, such as versioning, +dependency tracking, and collaboration. + +For example, if the asset manager supports a meta-version of "latest", +then the workflow inputs can be updated without having to edit the +workflow or move files around. + +Then, when the workflow completes, the output can be published back to +the asset manager, which typically creates a new version/revision +(rather than overwriting), and makes the output available for review and +for use by downstream tools. + +## Features + +- _OpenAssetIO Resolve Image_: An alternative to the built-in + _Load Image_ node that resolves an OpenAssetIO entity reference to an + image. + +- _OpenAssetIO Publish Image_: An alternative to the built-in _Save + Image_ node that publishes the output of a workflow to an OpenAssetIO + entity reference. + +## Requirements + +The plugin is known to work with + +- Python 3.11 +- [ComfyUI](https://comfy.org) 0.3.57 +- [OpenAssetIO](https://github.com/OpenAssetIO/OpenAssetIO) 1.0.0 +- [OpenAssetIO-MediaCreation](https://github.com/OpenAssetIO/OpenAssetIO-MediaCreation) + 1.0.0-alpha.12 + +## Installation + +Install [ComfyUI](https://docs.comfy.org/get_started). + +Clone this repository under `ComfyUI/custom_nodes`. + +From the repository root, install dependencies + +``` +pip install -r requirements.txt +``` + +Ensure the ComfyUI execution environment is configured correctly for an +OpenAssetIO host application. See the +[OpenAssetIO documentation](https://docs.openassetio.org/OpenAssetIO/runtime_configuration.html) +for general instructions on host application configuration. + +In particular, ensure the `OPENASSETIO_DEFAULT_CONFIG` environment +variable contains a path to a valid OpenAssetIO configuration file. + +## Development + +To install the dev dependencies and pre-commit (will run the +[Ruff](https://docs.astral.sh/ruff/) hook), from the repository root run + +```bash +pip install -e .[dev] +pre-commit install +``` + +> Note that installing this project to the Python environment has no +> effect on ComfyUI, since it loads plugins from the `custom_nodes` +> directory. However, installing the package helps with IDE code +> completion and linting; and of course ensures test/lint dependencies +> are installed. + +### Running Tests + +This project contains unit tests written in +[pytest](https://docs.pytest.org/en/stable/) in the `tests` directory. +To run the tests, from the repository root run + +```bash +pytest tests +``` + +### Linting + +The project makes use of the [Ruff](https://docs.astral.sh/ruff/) +linter, configured through the `pyproject.toml` file. To run Ruff, from +the repository root run + +```bash +ruff check . +``` + +## License + +Apache-2.0 - See [LICENSE](./LICENSE) file for details. + +## Contributing + +Please feel free to contribute pull requests or issues. Note that +contributions will require signing a CLA. + +See the OpenAssetIO contribution docs for how to structure +[commit messages](https://github.com/OpenAssetIO/OpenAssetIO/blob/main/doc/contributing/COMMITS.md), +the [pull request process](https://github.com/OpenAssetIO/OpenAssetIO/blob/main/doc/contributing/PULL_REQUESTS.md), +and [coding style guide](https://github.com/OpenAssetIO/OpenAssetIO/blob/main/doc/contributing/CODING_STYLE.md). diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..43745b6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,29 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +""" +Top-level package for openassetio-comfyui. + +This __init__.py will not be packaged, but is useful if the repository +is checked out under ComfyUI's custom_nodes directory, as it allows +ComfyUI to discover the nodes contained within. +""" + +import sys +import pathlib + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", +] + +__author__ = """Contributors to the OpenAssetIO project""" +__email__ = "openassetio-discussion@lists.aswf.io" +__version__ = "1.0.0" + +# Ensure src/ is on the path so we can import from there. +sys.path.append(str(pathlib.Path(__file__).parent / "src")) + +from openassetio_comfyui.nodes import NODE_CLASS_MAPPINGS +from openassetio_comfyui.nodes import NODE_DISPLAY_NAME_MAPPINGS diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..bb82c15 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,116 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "openassetio-comfyui" +version = "1.0.0" +description = "ComfyUI custom nodes adding asset resolution and publishing via OpenAssetIO" +authors = [ + { name = "Contributors to the OpenAssetIO project", email = "openassetio-discussion@lists.aswf.io" } +] +readme = "README.md" +license = { text = "Apache Software License 2.0" } +requires-python = ">=3.10,<3.12" +classifiers = [] +dependencies = [ + "openassetio>=1.0.0", + "openassetio-mediacreation>=1.0.0a12" +] + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # linting + "openassetio-manager-bal" # mock manager for testing +] + +[project.urls] +Repository = "https://github.com/OpenAssetIO/OpenAssetIO-ComfyUI" +BugTracker = "https://github.com/OpenAssetIO/OpenAssetIO-ComfyUI/issues" +Documentation = "https://github.com/OpenAssetIO/OpenAssetIO-ComfyUI/wiki" + + +[tool.comfy] +PublisherId = "OpenAssetIO" +DisplayName = "OpenAssetIO" +Icon = "" +Tags = [] +Repository = "https://github.com/OpenAssetIO/OpenAssetIO-ComfyUI" + +includes = [] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + +[tool.mypy] +files = "." + +# Use strict defaults +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +# Don't require test functions to include types +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +# extend-exclude = ["static", "ci/templates"] +line-length = 99 +src = ["src", "tests"] +target-version = "py311" + +# Add rules to ban exec/eval +[tool.ruff.lint] +select = [ + # See all rules here: https://docs.astral.sh/ruff/rules + "S102", # exec-builtin + "S307", # eval-used + "W293", + # The "F" series in Ruff stands for "Pyflakes" rules, which catch + # various Python syntax errors and undefined names. + "F", + # pydocstyle + "D", + # pycodestyle + "E", "W", + # pylint + "PL" +] + +ignore = [ + # "One-line docstring should fit on one line" - i.e. expects + # """summary""", not """\nsummary\n""", which is inconsistent with + # other OpenAssetIO projects. + "D200", + # "1 blank line required between summary line and description" - + # triggers if the summary spans more than one line. + "D205" +] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" + +[tool.ruff.lint.pydocstyle] +convention = "pep257" + +[tool.ruff.lint.pycodestyle] +max-doc-length = 72 +max-line-length = 99 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e3f7cd0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ + openassetio>=1.0.0 + openassetio-mediacreation>=1.0.0a12 diff --git a/src/openassetio_comfyui/__init__.py b/src/openassetio_comfyui/__init__.py new file mode 100644 index 0000000..d759700 --- /dev/null +++ b/src/openassetio_comfyui/__init__.py @@ -0,0 +1,17 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 + +"""Top-level package for openassetio-comfyui.""" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", +] + +__author__ = """Contributors to the OpenAssetIO project""" +__email__ = "openassetio-discussion@lists.aswf.io" +__version__ = "1.0.0" + +from .nodes import NODE_CLASS_MAPPINGS +from .nodes import NODE_DISPLAY_NAME_MAPPINGS diff --git a/src/openassetio_comfyui/nodes.py b/src/openassetio_comfyui/nodes.py new file mode 100644 index 0000000..7f759b5 --- /dev/null +++ b/src/openassetio_comfyui/nodes.py @@ -0,0 +1,417 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 +""" +OpenAssetIO nodes for ComfyUI. + +Also contains a singleton class representing the OpenAssetIO host +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, 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: + """ + Singleton class representing the OpenAssetIO host. + + Instantiates and exposes a manager and context, and provides common + utility methods. + + The OpenAssetIO manager must be configured via a config file, + located via the OPENASSETIO_DEFAULT_CONFIG environment variable. See + the OpenAssetIO documentation for more details. + """ + + _instance = None + + @classmethod + def instance(cls) -> "_OpenAssetIOHost": + """ + Get or create the singleton instance of the OpenAssetIO host. + """ + if cls._instance is None: + cls._instance = cls() + return cls._instance + + class _HostInterface(HostInterface): + def displayName(self) -> str: + return "ComfyUI" + + def identifier(self) -> str: + return "org.foundry.comfyui" + + def info(self) -> dict: + return super().info() + + class _LoggerInterface(LoggerInterface): + def __init__(self): + super().__init__() + self.__logger = logging.getLogger("openassetio-comfyui") + + def log(self, severity, message) -> None: + if severity == LoggerInterface.Severity.kDebug: + self.__logger.debug(message) + elif severity == LoggerInterface.Severity.kInfo: + self.__logger.info(message) + elif severity == LoggerInterface.Severity.kWarning: + self.__logger.warning(message) + elif severity == LoggerInterface.Severity.kError: + self.__logger.error(message) + elif severity == LoggerInterface.Severity.kCritical: + self.__logger.critical(message) + else: + self.__logger.log(logging.NOTSET, message) + + def __init__(self): + """ + Load and initialise an OpenAssetIO manager plugin. + """ + 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), + ], + 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: EntityReference | str, + access_mode: ResolveAccess = ResolveAccess.kRead, + ) -> str: + """ + Resolve an OpenAssetIO entity reference to a local file path. + + OpenAssetIO will raise exceptions if the entity reference is + invalid, if there is a problem resolving the entity, or if the + given entity location is not a `file://` URL. + + @throw RuntimeError if the entity reference resolves + successfully, but the entity has no location. + """ + if isinstance(entity_reference, str): + entity_reference = self.manager.createEntityReference(entity_reference) + + traits_data = self.manager.resolve( + entity_reference, {LocatableContentTrait.kId}, access_mode, self.context + ) + locatable_content_trait = LocatableContentTrait(traits_data) + url = locatable_content_trait.getLocation() + + if url is None: + raise RuntimeError( + f"Failed to resolve entity reference '{entity_reference}' to a location" + ) + + return self.file_url_path_converter.pathFromUrl(url) + + +class ResolveImage: + """ + Node to resolve an image from an OpenAssetIO entity reference. + + The non-OpenAssetIO logic is largely duplicated from the built-in + ComfyUI LoadImage node (at least as of 4f5812b9). + """ + + # Tooltip to display when hovering over the node. + DESCRIPTION = "Resolve images from an asset manager." + # Menu category. + CATEGORY = "image" + # Function to call when node is executed. + FUNCTION = "resolve_image" + # Node outputs. + RETURN_TYPES = ("IMAGE", "MASK") + + @classmethod + def INPUT_TYPES(cls) -> dict: + """ + Input sockets and widgets. + """ + return { + "required": { + "entity_reference": ( + "STRING", + { + "multiline": False, + "default": "", + "placeholder": "Entity reference...", + "tooltip": "The entity to resolve", + }, + ) + } + } + + @classmethod + def IS_CHANGED(cls, entity_reference: str) -> str: + """ + Resolve the entity reference to a file path, then return a hash + of the file contents. + + This allows ComfyUI to determine if the node needs to be + re-executed. + + Non-OpenAssetIO logic largely duplicated from the built-in + ComfyUI LoadImage node (at least as of 4f5812b9). + """ + image_path = _OpenAssetIOHost.instance().resolve_to_path(entity_reference) + m = hashlib.sha256() + with open(image_path, "rb") as f: + m.update(f.read()) + return m.digest().hex() + + @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 resolve_image(self, entity_reference: str) -> tuple[torch.Tensor, torch.Tensor]: + """ + Resolve an OpenAssetIO entity reference to a file path, then + load the image(s) at that path and return as tensors. + + Non-OpenAssetIO logic largely duplicated from the built-in + ComfyUI LoadImage node (at least as of 4f5812b9). + """ + image_path = _OpenAssetIOHost.instance().resolve_to_path(entity_reference) + + sequence = node_helpers.pillow(Image.open, image_path) + + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ["MPO"] + + for frame in ImageSequence.Iterator(sequence): + frame = node_helpers.pillow(ImageOps.exif_transpose, frame) # noqa: PLW2901 - frame overwrite + + if frame.mode == "I": + frame = frame.point(lambda px: px * (1 / 255)) # noqa: PLW2901 - frame overwrite + image = frame.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if "A" in frame.getbands(): + mask = np.array(frame.getchannel("A")).astype(np.float32) / 255.0 + mask = 1.0 - torch.from_numpy(mask) + elif frame.mode == "P" and "transparency" in frame.info: + mask = np.array(frame.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0 + mask = 1.0 - torch.from_numpy(mask) + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + output_images.append(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and sequence.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + + 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/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..5b381ee --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,6 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 +""" +Unit test package for openassetio-comfyui. +""" diff --git a/tests/resources/bal_db.json b/tests/resources/bal_db.json new file mode 100644 index 0000000..d443d3b --- /dev/null +++ b/tests/resources/bal_db.json @@ -0,0 +1,65 @@ +{ + "managementPolicy": { + "read": { + "default": { + "openassetio-mediacreation:managementPolicy.Managed": {} + } + }, + "write": { + "default": { + "openassetio-mediacreation:managementPolicy.Managed": {} + } + }, + "managerDriven": { + "default": { + "openassetio-mediacreation:managementPolicy.Managed": {} + } + } + }, + "entities": { + "image_file": { + "versions": [ + { + "traits": { + "openassetio-mediacreation:usage.Entity": {}, + "openassetio-mediacreation:twoDimensional.Image": {}, + "openassetio-mediacreation:content.LocatableContent": { + "location": "file://${test_tmp_dir}/cat.png" + } + } + } + ] + }, + "no_file": { + "versions": [ + { + "traits": { + "openassetio-mediacreation:usage.Entity": {}, + "openassetio-mediacreation:twoDimensional.Image": {} + } + } + ] + }, + "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/resources/openassetio_config.toml b/tests/resources/openassetio_config.toml new file mode 100644 index 0000000..dc9d08c --- /dev/null +++ b/tests/resources/openassetio_config.toml @@ -0,0 +1,10 @@ +# openassetio-comfyui +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The Foundry Visionmongers Ltd + +[manager] +identifier = "org.openassetio.examples.manager.bal" + +[manager.settings] +library_path = "${config_dir}/bal_db.json" +simulated_query_latency_ms = 0 diff --git a/tests/test_openassetio_comfyui.py b/tests/test_openassetio_comfyui.py new file mode 100644 index 0000000..6eb8b4e --- /dev/null +++ b/tests/test_openassetio_comfyui.py @@ -0,0 +1,349 @@ +# openassetio-comfyui +# Copyright (c) 2025 The Foundry Visionmongers Ltd +# SPDX-License-Identifier: Apache-2.0 +""" +Tests for `openassetio-comfyui` package. +""" + +import contextlib +import hashlib +# D101: "Missing docstring in public class" +# D102: "Missing docstring in public method" +# ruff: noqa: D101,D102 + +import inspect +import pathlib + +import numpy as np +import pytest +import torch +from PIL import Image + +from openassetio_comfyui.nodes import ResolveImage, _OpenAssetIOHost, PublishImage + +import folder_paths + + +class Test_ResolveImage_init: + """ + Test that the node can be instantiated. + """ + + def test_has_correct_type(self, resolve_image_node): + assert isinstance(resolve_image_node, ResolveImage) + + +class Test_ResolveImage_constants: + """ + Test the node's metadata. + """ + + def test_has_expected_values(self): + assert ResolveImage.DESCRIPTION == "Resolve images from an asset manager." + assert ResolveImage.CATEGORY == "image" + assert ResolveImage.FUNCTION == "resolve_image" + assert ResolveImage.RETURN_TYPES == ("IMAGE", "MASK") + + def test_function_matches(self): + assert inspect.isfunction(getattr(ResolveImage, ResolveImage.FUNCTION)) + + +class Test_ResolveImage_INPUT_TYPES: + def test_has_expected_structure(self): + input_types = ResolveImage.INPUT_TYPES() + + assert input_types == { + "required": { + "entity_reference": ( + "STRING", + { + "multiline": False, + "placeholder": "Entity reference...", + "default": "", + "tooltip": "The entity to resolve", + }, + ) + } + } + + +class Test_ResolveImage_VALIDATE_INPUTS: + def test_when_is_a_reference_then_returns_true(self): + valid = ResolveImage.VALIDATE_INPUTS("bal:///") + assert valid + + def test_when_not_a_reference_then_returns_false(self): + valid = ResolveImage.VALIDATE_INPUTS("notbal:///") + assert not valid + + def test_when_no_manager_found_then_raises(self, assert_raises_missing_manager): + with assert_raises_missing_manager(): + ResolveImage.VALIDATE_INPUTS("bal:///") + + +class Test_ResolveImage_IS_CHANGED: + def test_when_reference_resolves_to_file_then_hash_matches_file(self, image_file_hash): + actual_file_hash = ResolveImage.IS_CHANGED("bal:///image_file") + + assert actual_file_hash == image_file_hash + + def test_when_reference_has_no_path_then_raises(self, image_file_hash): + expected_error_message = ( + "Failed to resolve entity reference 'bal:///no_file' to a location" + ) + + with pytest.raises(RuntimeError, match=expected_error_message): + ResolveImage.IS_CHANGED("bal:///no_file") + + def test_when_no_manager_found_then_raises(self, assert_raises_missing_manager): + with assert_raises_missing_manager(): + ResolveImage.IS_CHANGED("bal:///image_file") + + +class Test_ResolveImage_resolve_image: + def test_when_reference_resolves_to_file_then_returns_image( + self, resolve_image_node, image_file_path + ): + images, masks = resolve_image_node.resolve_image("bal:///image_file") + + assert isinstance(images, torch.Tensor) + assert images.shape == (1, 2, 2, 3) + assert isinstance(masks, torch.Tensor) + assert masks.shape == (1, 2, 2) + + image = images[0] + mask = masks[0] + + # Assert channels at each pixel. + assert torch.equal(image[0, 0], torch.tensor([0, 0, 0])) + assert torch.equal(image[0, 1], torch.tensor([1, 1, 1])) + assert torch.equal(image[1, 0], torch.tensor([0, 0, 0])) + assert torch.equal(image[1, 1], torch.tensor([1, 1, 1])) + + # Note: mask is `1 - alpha`. + assert mask[0, 0].item() == 0 + assert mask[0, 1].item() == 1 + assert mask[1, 0].item() == 1 + assert mask[1, 1].item() == 0 + + def test_when_no_manager_found_then_raises( + self, resolve_image_node, assert_raises_missing_manager + ): + with assert_raises_missing_manager(): + 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): + """ + Fixture to assert that a RuntimeError is raised when no manager can + be created. + + Since a lazily-created singleton is used to manage the OpenAssetIO + host, this assertion must be tested for every function that uses + that singleton. + """ + + @contextlib.contextmanager + def assert_ctx(): + monkeypatch.delenv("OPENASSETIO_DEFAULT_CONFIG", raising=False) + expected_error_msg = ( + "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." + ) + with pytest.raises(RuntimeError, match=expected_error_msg): + yield + + 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): + """ + Fixture to provide the hash of the temporary image file. + """ + hasher = hashlib.sha256() + hasher.update(image_file_path.read_bytes()) + return hasher.hexdigest() + + +@pytest.fixture +def image_file_path(monkeypatch, tmp_path): + """ + Fixture to provide a temporary image file. + """ + # See bal_db.json - the resolved path will interpolate this env var. + monkeypatch.setenv("test_tmp_dir", str(tmp_path)) + # See bal_db.json - the resolved path for bal:///cat + file_path = pathlib.Path(tmp_path) / "cat.png" + + content = np.array( + [[[0, 0, 0, 255], [255, 255, 255, 0]], [[0, 0, 0, 0], [255, 255, 255, 255]]] + ).astype(np.uint8) + img = Image.fromarray(content) + img.save(file_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(): + """ + Fixture to create an ResolveImage node instance. + """ + 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): + """ + Fixture to set the OPENASSETIO_DEFAULT_CONFIG environment variable + to point to a test config file. + """ + test_config_path = resources_dir / "openassetio_config.toml" + monkeypatch.setenv("OPENASSETIO_DEFAULT_CONFIG", str(test_config_path)) + + +@pytest.fixture(autouse=True) +def reset_singleton_host(monkeypatch): + """ + Fixture to reset the singleton instance before each test. + """ + _OpenAssetIOHost._instance = None + + +@pytest.fixture(scope="module") +def resources_dir(): + """ + Fixture to provide the path to the resources directory. + """ + return pathlib.Path(__file__).parent / "resources"