From e8dc58f38a52c3dd2bc1635a001d0c7708641f49 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Sat, 28 Oct 2023 23:31:49 -0400
Subject: [PATCH] Convert more node defs to ComfyTypedNodeBase

---
 .../Api/Comfy/Nodes/ComfyNodeBuilder.cs       | 52 +++++++++++++------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
index 52aa272e..767af550 100644
--- a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
+++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
@@ -230,13 +230,9 @@ public class ComfyNodeBuilder
         };
     }
 
-    public static NamedComfyNode<VAENodeConnection> VAELoader(string name, string vaeModelName)
+    public record VAELoader : ComfyTypedNodeBase<VAENodeConnection>
     {
-        return new NamedComfyNode<VAENodeConnection>(name)
-        {
-            ClassType = "VAELoader",
-            Inputs = new Dictionary<string, object?> { ["vae_name"] = vaeModelName }
-        };
+        public required string VaeName { get; init; }
     }
 
     public static NamedComfyNode<ModelNodeConnection, ClipNodeConnection> LoraLoader(
@@ -262,16 +258,10 @@ public class ComfyNodeBuilder
         };
     }
 
-    public static NamedComfyNode<ModelNodeConnection> CheckpointLoaderSimple(
-        string name,
-        string modelName
-    )
+    public record CheckpointLoaderSimple
+        : ComfyTypedNodeBase<ModelNodeConnection, ClipNodeConnection, VAENodeConnection>
     {
-        return new NamedComfyNode<ModelNodeConnection>(name)
-        {
-            ClassType = "CheckpointLoaderSimple",
-            Inputs = new Dictionary<string, object?> { ["ckpt_name"] = modelName }
-        };
+        public required string CkptName { get; init; }
     }
 
     public static NamedComfyNode<ModelNodeConnection> FreeU(
@@ -739,6 +729,22 @@ public class ComfyNodeBuilder
         return primary.Match(latent => latent, image => Lambda_ImageToLatent(image, vae));
     }
 
+    /// <summary>
+    /// Get or convert latest primary connection to latent
+    /// </summary>
+    public LatentNodeConnection GetPrimaryAsLatent(VAENodeConnection vae)
+    {
+        if (Connections.Primary?.IsT0 == true)
+        {
+            return Connections.Primary.AsT0;
+        }
+
+        return GetPrimaryAsLatent(
+            Connections.Primary ?? throw new NullReferenceException("No primary connection"),
+            vae
+        );
+    }
+
     /// <summary>
     /// Get or convert latest primary connection to image
     /// </summary>
@@ -766,6 +772,22 @@ public class ComfyNodeBuilder
         return primary.Match(latent => Lambda_LatentToImage(latent, vae), image => image);
     }
 
+    /// <summary>
+    /// Get or convert latest primary connection to image
+    /// </summary>
+    public ImageNodeConnection GetPrimaryAsImage(VAENodeConnection vae)
+    {
+        if (Connections.Primary?.IsT1 == true)
+        {
+            return Connections.Primary.AsT1;
+        }
+
+        return GetPrimaryAsImage(
+            Connections.Primary ?? throw new NullReferenceException("No primary connection"),
+            vae
+        );
+    }
+
     /// <summary>
     /// Convert to a NodeDictionary
     /// </summary>