From d12e04a69ba249668a57730d3ae6061aaade1017 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Wed, 21 Feb 2024 16:08:30 -0500
Subject: [PATCH] Add NodeConnectionBase Json converter

---
 .../Json/NodeConnectionBaseJsonConverter.cs   | 49 +++++++++++++++++++
 StabilityMatrix.Core/Inference/ComfyClient.cs |  6 ++-
 .../Api/Comfy/NodeTypes/NodeConnectionBase.cs |  6 ++-
 3 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 StabilityMatrix.Core/Converters/Json/NodeConnectionBaseJsonConverter.cs

diff --git a/StabilityMatrix.Core/Converters/Json/NodeConnectionBaseJsonConverter.cs b/StabilityMatrix.Core/Converters/Json/NodeConnectionBaseJsonConverter.cs
new file mode 100644
index 00000000..1d90858e
--- /dev/null
+++ b/StabilityMatrix.Core/Converters/Json/NodeConnectionBaseJsonConverter.cs
@@ -0,0 +1,49 @@
+using System.Reflection;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+
+namespace StabilityMatrix.Core.Converters.Json;
+
+public class NodeConnectionBaseJsonConverter : JsonConverter<NodeConnectionBase>
+{
+    /// <inheritdoc />
+    public override NodeConnectionBase Read(
+        ref Utf8JsonReader reader,
+        Type typeToConvert,
+        JsonSerializerOptions options
+    )
+    {
+        // Read as Data array
+        reader.Read();
+        var data = new object[2];
+        reader.Read();
+        data[0] = reader.GetString() ?? throw new JsonException("Expected string for node name");
+        reader.Read();
+        data[1] = reader.GetInt32();
+        reader.Read();
+
+        if (Activator.CreateInstance(typeToConvert) is not NodeConnectionBase instance)
+        {
+            throw new JsonException($"Failed to create instance of {typeToConvert}");
+        }
+
+        var propertyInfo =
+            typeToConvert.GetProperty("Data", BindingFlags.Public | BindingFlags.Instance)
+            ?? throw new JsonException($"Failed to get Data property of {typeToConvert}");
+
+        propertyInfo.SetValue(instance, data);
+
+        return instance;
+    }
+
+    /// <inheritdoc />
+    public override void Write(Utf8JsonWriter writer, NodeConnectionBase value, JsonSerializerOptions options)
+    {
+        // Write as Data array
+        writer.WriteStartArray();
+        writer.WriteStringValue(value.Data?[0] as string);
+        writer.WriteNumberValue((int)value.Data?[1]!);
+        writer.WriteEndArray();
+    }
+}
diff --git a/StabilityMatrix.Core/Inference/ComfyClient.cs b/StabilityMatrix.Core/Inference/ComfyClient.cs
index 248396f0..72ae9a27 100644
--- a/StabilityMatrix.Core/Inference/ComfyClient.cs
+++ b/StabilityMatrix.Core/Inference/ComfyClient.cs
@@ -33,7 +33,11 @@ public class ComfyClient : InferenceClientBase
         new()
         {
             PropertyNamingPolicy = JsonNamingPolicies.SnakeCaseLower,
-            Converters = { new OneOfJsonConverter<string, StringNodeConnection>() }
+            Converters =
+            {
+                new NodeConnectionBaseJsonConverter(),
+                new OneOfJsonConverter<string, StringNodeConnection>()
+            }
         };
 
     // ReSharper disable once MemberCanBePrivate.Global
diff --git a/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnectionBase.cs b/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnectionBase.cs
index 01bbff91..19dd1ad5 100644
--- a/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnectionBase.cs
+++ b/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnectionBase.cs
@@ -1,5 +1,9 @@
-namespace StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+using System.Text.Json.Serialization;
+using StabilityMatrix.Core.Converters.Json;
 
+namespace StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+
+[JsonConverter(typeof(NodeConnectionBaseJsonConverter))]
 public abstract class NodeConnectionBase
 {
     /// <summary>