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