From 3653fb475e38a17dd5162d86edbd963ae67e788b Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Fri, 1 Mar 2024 16:43:27 -0500
Subject: [PATCH] Fix missing node causing connection error

---
 StabilityMatrix.Core/Inference/ComfyClient.cs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/StabilityMatrix.Core/Inference/ComfyClient.cs b/StabilityMatrix.Core/Inference/ComfyClient.cs
index 9e7ab6bf..3f3b894e 100644
--- a/StabilityMatrix.Core/Inference/ComfyClient.cs
+++ b/StabilityMatrix.Core/Inference/ComfyClient.cs
@@ -471,8 +471,9 @@ public class ComfyClient : InferenceClientBase
     {
         var response = await comfyApi.GetObjectInfo(nodeName, cancellationToken).ConfigureAwait(false);
 
-        var info = response[nodeName];
-        return info.Input.GetOptionalValueAsNestedList(optionName);
+        var info = response.GetValueOrDefault(nodeName);
+
+        return info?.Input.GetOptionalValueAsNestedList(optionName);
     }
 
     protected override void Dispose(bool disposing)