Add OneOfJsonConverter for ComfyClient

StabilityMatrix.Core/Converters/Json/OneOfJsonConverter.cs

@@ -0,0 +1,63 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using OneOf;
+
+namespace StabilityMatrix.Core.Converters.Json;
+
+public class OneOfJsonConverter<T1, T2> : JsonConverter<OneOf<T1, T2>>
+{
+    /// <inheritdoc />
+    public override OneOf<T1, T2> Read(
+        ref Utf8JsonReader reader,
+        Type typeToConvert,
+        JsonSerializerOptions options
+    )
+    {
+        // Not sure how else to do this without polymorphic type markers but that would not serialize into T1/T2
+        // So just try to deserialize T1, if it fails, try T2
+        Exception? t1Exception = null;
+        Exception? t2Exception = null;
+
+        try
+        {
+            if (JsonSerializer.Deserialize<T1>(ref reader, options) is { } t1)
+            {
+                return t1;
+            }
+        }
+        catch (JsonException e)
+        {
+            t1Exception = e;
+        }
+
+        try
+        {
+            if (JsonSerializer.Deserialize<T2>(ref reader, options) is { } t2)
+            {
+                return t2;
+            }
+        }
+        catch (JsonException e)
+        {
+            t2Exception = e;
+        }
+
+        throw new JsonException(
+            $"Failed to deserialize OneOf<{typeof(T1)}, {typeof(T2)}> as either {typeof(T1)} or {typeof(T2)}",
+            new AggregateException([t1Exception, t2Exception])
+        );
+    }
+
+    /// <inheritdoc />
+    public override void Write(Utf8JsonWriter writer, OneOf<T1, T2> value, JsonSerializerOptions options)
+    {
+        if (value.IsT0)
+        {
+            JsonSerializer.Serialize(writer, value.AsT0, options);
+        }
+        else
+        {
+            JsonSerializer.Serialize(writer, value.AsT1, options);
+        }
+    }
+}

StabilityMatrix.Core/Inference/ComfyClient.cs

@@ -6,11 +6,13 @@ using NLog;
 using Polly.Contrib.WaitAndRetry;
 using Refit;
 using StabilityMatrix.Core.Api;
+using StabilityMatrix.Core.Converters.Json;
 using StabilityMatrix.Core.Exceptions;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.Api.Comfy;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
+using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
 using StabilityMatrix.Core.Models.Api.Comfy.WebSocketData;
 using StabilityMatrix.Core.Models.FileInterfaces;
 using Websocket.Client;
@@ -27,8 +29,12 @@ public class ComfyClient : InferenceClientBase
     private readonly IComfyApi comfyApi;
     private bool isDisposed;
 
-    private JsonSerializerOptions jsonSerializerOptions =
-        new() { PropertyNamingPolicy = JsonNamingPolicies.SnakeCaseLower, };
+    private readonly JsonSerializerOptions jsonSerializerOptions =
+        new()
+        {
+            PropertyNamingPolicy = JsonNamingPolicies.SnakeCaseLower,
+            Converters = { new OneOfJsonConverter<string, StringNodeConnection>() }
+        };
 
     // ReSharper disable once MemberCanBePrivate.Global
     public string ClientId { get; } = Guid.NewGuid().ToString();
@@ -87,7 +93,13 @@ public class ComfyClient : InferenceClientBase
 
     public ComfyClient(IApiFactory apiFactory, Uri baseAddress)
     {
-        comfyApi = apiFactory.CreateRefitClient<IComfyApi>(baseAddress);
+        comfyApi = apiFactory.CreateRefitClient<IComfyApi>(
+            baseAddress,
+            new RefitSettings
+            {
+                ContentSerializer = new SystemTextJsonContentSerializer(jsonSerializerOptions),
+            }
+        );
         BaseAddress = baseAddress;
 
         // Setup websocket client