From 479ceafc2e9a1f196772bb80e226d10f8dd4eab3 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Sun, 6 Aug 2023 17:39:40 -0400
Subject: [PATCH] Add PackagePair overload for ConnectAsync

---
 .../DesignData/MockInferenceClientManager.cs  |  7 +++++
 .../Services/IInferenceClientManager.cs       |  3 ++
 .../Services/InferenceClientManager.cs        | 29 +++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/StabilityMatrix.Avalonia/DesignData/MockInferenceClientManager.cs b/StabilityMatrix.Avalonia/DesignData/MockInferenceClientManager.cs
index 17a94b14..2e1e6658 100644
--- a/StabilityMatrix.Avalonia/DesignData/MockInferenceClientManager.cs
+++ b/StabilityMatrix.Avalonia/DesignData/MockInferenceClientManager.cs
@@ -4,6 +4,7 @@ using System.Threading.Tasks;
 using CommunityToolkit.Mvvm.ComponentModel;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Core.Inference;
+using StabilityMatrix.Core.Models;
 
 namespace StabilityMatrix.Avalonia.DesignData;
 
@@ -30,6 +31,12 @@ public class MockInferenceClientManager : ObservableObject, IInferenceClientMana
         return Task.CompletedTask;
     }
 
+    /// <inheritdoc />
+    public Task ConnectAsync(PackagePair packagePair)
+    {
+        return Task.CompletedTask;
+    }
+
     public Task CloseAsync()
     {
         return Task.CompletedTask;
diff --git a/StabilityMatrix.Avalonia/Services/IInferenceClientManager.cs b/StabilityMatrix.Avalonia/Services/IInferenceClientManager.cs
index 25020edc..d869c675 100644
--- a/StabilityMatrix.Avalonia/Services/IInferenceClientManager.cs
+++ b/StabilityMatrix.Avalonia/Services/IInferenceClientManager.cs
@@ -4,6 +4,7 @@ using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading.Tasks;
 using StabilityMatrix.Core.Inference;
+using StabilityMatrix.Core.Models;
 
 namespace StabilityMatrix.Avalonia.Services;
 
@@ -19,5 +20,7 @@ public interface IInferenceClientManager : IDisposable,  INotifyPropertyChanged,
     
     Task ConnectAsync();
 
+    Task ConnectAsync(PackagePair packagePair);
+
     Task CloseAsync();
 }
diff --git a/StabilityMatrix.Avalonia/Services/InferenceClientManager.cs b/StabilityMatrix.Avalonia/Services/InferenceClientManager.cs
index f643fe1a..8debd337 100644
--- a/StabilityMatrix.Avalonia/Services/InferenceClientManager.cs
+++ b/StabilityMatrix.Avalonia/Services/InferenceClientManager.cs
@@ -3,8 +3,13 @@ using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading.Tasks;
 using CommunityToolkit.Mvvm.ComponentModel;
+using Sentry;
 using StabilityMatrix.Core.Api;
+using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Inference;
+using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.FileInterfaces;
+using StabilityMatrix.Core.Models.Packages;
 
 namespace StabilityMatrix.Avalonia.Services;
 
@@ -16,6 +21,8 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
 {
     private readonly IApiFactory apiFactory;
     
+    // Current 
+    
     [ObservableProperty, NotifyPropertyChangedFor(nameof(IsConnected))]
     private ComfyClient? client;
     
@@ -57,6 +64,28 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
         await LoadSharedPropertiesAsync();
     }
     
+    public async Task ConnectAsync(PackagePair packagePair)
+    {
+        if (IsConnected) return;
+
+        if (packagePair.BasePackage is not ComfyUI)
+        {
+            throw new ArgumentException("Base package is not ComfyUI", nameof(packagePair));
+        }
+
+        var tempClient = new ComfyClient(apiFactory, new Uri("http://127.0.0.1:8188"));
+
+        // Add output dir if available
+        if (packagePair.InstalledPackage.FullPath is { } path)
+        {
+            tempClient.OutputImagesDir = new DirectoryPath(path, "output");
+        }
+        
+        await tempClient.ConnectAsync();
+        Client = tempClient;
+        await LoadSharedPropertiesAsync();
+    }
+    
     public async Task CloseAsync()
     {
         if (!IsConnected) return;