From ff26180e8bd69ba562583dd9305c031c572889c1 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Thu, 4 Apr 2024 23:08:05 -0400
Subject: [PATCH] Apply post install dep updates for ComfyUI ext update

---
 .../Models/Packages/ComfyUI.cs                | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
index 19843e75..c620182c 100644
--- a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
+++ b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
@@ -511,6 +511,32 @@ public class ComfyUI(
             }
         }
 
+        /// <inheritdoc />
+        public override async Task UpdateExtensionAsync(
+            InstalledPackageExtension installedExtension,
+            InstalledPackage installedPackage,
+            PackageExtensionVersion? version = null,
+            IProgress<ProgressReport>? progress = null,
+            CancellationToken cancellationToken = default
+        )
+        {
+            await base.UpdateExtensionAsync(
+                installedExtension,
+                installedPackage,
+                version,
+                progress,
+                cancellationToken
+            )
+                .ConfigureAwait(false);
+
+            cancellationToken.ThrowIfCancellationRequested();
+
+            var installedDirs = installedExtension.Paths.OfType<DirectoryPath>().Where(dir => dir.Exists);
+
+            await PostInstallAsync(installedPackage, installedDirs, progress, cancellationToken)
+                .ConfigureAwait(false);
+        }
+
         /// <inheritdoc />
         public override async Task InstallExtensionAsync(
             PackageExtension extension,
@@ -539,6 +565,20 @@ public class ComfyUI(
                 .Select(path => cloneRoot.JoinDir(path!))
                 .Where(dir => dir.Exists);
 
+            await PostInstallAsync(installedPackage, installedDirs, progress, cancellationToken)
+                .ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// Runs post install / update tasks (i.e. install.py, requirements.txt)
+        /// </summary>
+        private async Task PostInstallAsync(
+            InstalledPackage installedPackage,
+            IEnumerable<DirectoryPath> installedDirs,
+            IProgress<ProgressReport>? progress = null,
+            CancellationToken cancellationToken = default
+        )
+        {
             foreach (var installedDir in installedDirs)
             {
                 cancellationToken.ThrowIfCancellationRequested();