Merge pull request #1 from feltech/work/resolve_and_publish_images
Add ResolveImage and PublishImage nodes
This commit is contained in:
David Feltell
@@ -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
|
||||
@@ -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.
|
||||
```
|
||||
@@ -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 .
|
||||
@@ -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 }}
|
||||
@@ -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
|
||||
+118
@@ -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/
|
||||
@@ -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
|
||||
+13
@@ -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
|
||||
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
openassetio-comfyui
|
||||
Copyright (c) 2025 The Foundry Visionmongers Ltd.
|
||||
|
||||
This product includes software developed at
|
||||
The Foundry Visionmongers Ltd.
|
||||
@@ -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).
|
||||
+29
@@ -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
|
||||
+116
@@ -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
|
||||
@@ -0,0 +1,2 @@
|
||||
openassetio>=1.0.0
|
||||
openassetio-mediacreation>=1.0.0a12
|
||||
@@ -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
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user