From 7467fa5e57467da1e7afdf7aee6feda074f3d834 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 19 May 2023 16:10:16 +1000
Subject: [PATCH] fix(ui): fix node ui type hints

---
 invokeai/app/invocations/baseinvocation.py |  1 +
 invokeai/app/invocations/latent.py         | 16 +++-------------
 2 files changed, 4 insertions(+), 13 deletions(-)

diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py
index 5f75cfd3b8..4ce3e839b6 100644
--- a/invokeai/app/invocations/baseinvocation.py
+++ b/invokeai/app/invocations/baseinvocation.py
@@ -96,6 +96,7 @@ class UIConfig(TypedDict, total=False):
             "image",
             "latents",
             "model",
+            "control",
         ],
     ]
     tags: List[str]
diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py
index f9c9f2b30f..1befe483f0 100644
--- a/invokeai/app/invocations/latent.py
+++ b/invokeai/app/invocations/latent.py
@@ -186,9 +186,10 @@ class TextToLatentsInvocation(BaseInvocation):
     class Config(InvocationConfig):
         schema_extra = {
             "ui": {
-                "tags": ["latents", "image"],
+                "tags": ["latents"],
                 "type_hints": {
-                  "model": "model"
+                  "model": "model",
+                  "control": "control",
                 }
             },
         }
@@ -355,17 +356,6 @@ class LatentsToLatentsInvocation(TextToLatentsInvocation):
     latents: Optional[LatentsField] = Field(description="The latents to use as a base image")
     strength: float = Field(default=0.5, description="The strength of the latents to use")
 
-    # Schema customisation
-    class Config(InvocationConfig):
-        schema_extra = {
-            "ui": {
-                "tags": ["latents"],
-                "type_hints": {
-                    "model": "model"
-                }
-            },
-        }
-
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         noise = context.services.latents.get(self.noise.latents_name)
         latent = context.services.latents.get(self.latents.latents_name)