diff --git a/invokeai/app/invocations/primitives.py b/invokeai/app/invocations/primitives.py index 587cbaef51..598fbac8a3 100644 --- a/invokeai/app/invocations/primitives.py +++ b/invokeai/app/invocations/primitives.py @@ -1,6 +1,6 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) -from typing import Optional +from typing import Any, Optional import torch @@ -26,6 +26,7 @@ from invokeai.app.invocations.fields import ( SD3ConditioningField, TensorField, UIComponent, + UIType, ) from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.shared.invocation_context import InvocationContext @@ -535,3 +536,25 @@ class BoundingBoxInvocation(BaseInvocation): # endregion + + +@invocation_output("any_output") +class AnyOutput(BaseInvocationOutput): + value: Any = OutputField(description="The output value", ui_type=UIType.Any) + + +@invocation( + "switcher", + title="Switcher", + tags=["primitives", "switcher"], + category="primitives", + version="1.0.0", +) +class SwitcherInvocation(BaseInvocation): + a: Any = InputField(description="The first input", ui_type=UIType.Any) + b: Any = InputField(description="The second input", ui_type=UIType.Any) + switch: bool = InputField(description="Switch between the two inputs") + + def invoke(self, context: InvocationContext) -> AnyOutput: + value = self.a if self.switch else self.b + return AnyOutput(value=value)