From 2cda6cbf14497f0358e9e84a7b5c52b8d689ef88 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 29 Apr 2024 09:12:33 +1000 Subject: [PATCH] feat(nodes): attempt to look up invoke return types by name --- invokeai/app/invocations/baseinvocation.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py index d7902c0397..ab3cb9fc43 100644 --- a/invokeai/app/invocations/baseinvocation.py +++ b/invokeai/app/invocations/baseinvocation.py @@ -4,6 +4,7 @@ from __future__ import annotations import inspect import re +import sys import warnings from abc import ABC, abstractmethod from enum import Enum @@ -481,7 +482,11 @@ def invocation( # And validate that `invoke()` returns a subclass of `BaseInvocationOutput invoke_return_annotation = signature(cls.invoke).return_annotation + try: + if isinstance(invoke_return_annotation, str): + invoke_return_annotation = getattr(sys.modules[cls.__module__], invoke_return_annotation) + assert invoke_return_annotation is not BaseInvocationOutput # TODO(psyche): If `invoke()` is not defined, `return_annotation` ends up as the string # "BaseInvocationOutput". This may be a pydantic bug: https://github.com/pydantic/pydantic/issues/7978