experimental first pass comfy torch version updatinator

Author: mcmonkey4eva

src/BuiltinExtensions/ComfyUIBackend/Assets/comfy_workflow_editor_helper.js

@@ -1,6 +1,72 @@
 /** If true, the workflow iframe is present. If false, the tab has never been opened, or loading failed. */
 let hasComfyLoaded = false;
 
+/** Helper class managing the "ComfyUI Torch Versions" card on the Server Info tab. */
+class ComfyTorchManager {
+
+    constructor() {
+        getRequiredElementById('serverinfotabbutton').addEventListener('click', () => this.refresh());
+        getRequiredElementById('servertabbutton').addEventListener('click', () => this.refresh());
+        let serverInfoTab = document.getElementById('Server-Info');
+        let collection = createDiv(null, 'card-collection-inline');
+        this.card = createDiv('comfy_torch_card', 'card border-secondary mb-3 card-center-container');
+        this.card.dataset.requiredpermission = 'view_backends_list';
+        this.card.style.display = 'none';
+        this.card.innerHTML = `<div class="card-header translate">ComfyUI Torch Versions</div><div class="card-body"><span class="card-text" id="comfy_torch_card_body">(Loading...)</span></div>`;
+        collection.appendChild(this.card);
+        serverInfoTab.appendChild(collection);
+        this.bodyElem = getRequiredElementById('comfy_torch_card_body');
+    }
+
+    /** Refreshes the torch install list from the server and rebuilds the card body. */
+    refresh() {
+        if (!permissions.hasPermission('view_backends_list')) {
+            return;
+        }
+        genericRequest('ComfyListTorchInstalls', {}, (data) => {
+            if (!data.installs || data.installs.length == 0) {
+                this.card.style.display = 'none';
+                return;
+            }
+            this.card.style.display = '';
+            let html = `Installed Comfy Torch Version(s):<br><br>`
+                + `<table class="simple-table"><tr><th>Install Folder</th><th>Torch Version</th><th>Backend IDs</th><th>Action</th></tr>`;
+            for (let install of data.installs) {
+                let action;
+                if (install.can_update) {
+                    action = `<button class="basic-button translate" onclick="comfyTorchManager.updateTorch(this, '${escapeHtml(install.path)}', ${install.backend_ids[0]})">Update Torch</button>`;
+                }
+                else {
+                    action = `(None)`;
+                }
+                html += `<tr><td><code>${escapeHtml(install.path)}</code></td><td><code>${escapeHtml(install.torch_version)}</code></td><td>${escapeHtml(install.backend_ids.join(', '))}</td><td>${action}</td></tr>`;
+            }
+            html += `</table>`;
+            this.bodyElem.innerHTML = html;
+        });
+    }
+
+    /** Triggered by the Update Torch button, to run a torch update for one install folder. */
+    updateTorch(button, path, backendId) {
+        if (!confirm(`Are you sure you want to update PyTorch for the ComfyUI install at:\n${path}\n\nThis is experimental! It will take a while, download several gigabytes of data, and might even break things!`)) {
+            return;
+        }
+        button.disabled = true;
+        button.parentElement.querySelectorAll('.installing_info').forEach(e => e.remove());
+        let status = createDiv(null, 'installing_info', 'Updating torch (this may take several minutes, check server logs for details)...');
+        button.parentElement.appendChild(status);
+        genericRequest('ComfyUpdateTorch', { 'backendId': backendId }, (data) => {
+            status.innerText = 'Torch updated, backends restarting.';
+            setTimeout(() => this.refresh(), 3000);
+        }, 0, (e) => {
+            status.innerText = 'Failed to update torch: ' + e;
+            button.disabled = false;
+        });
+    }
+}
+
+let comfyTorchManager = new ComfyTorchManager();
+
 /** Helper class for managing the Comfy workflow tab. */
 class ComfyWorkflowHelpers {
 
@@ -1357,6 +1423,9 @@ featureSetChangedCallbacks.push(() => {
         comfyHasTriedToLoad = true;
         comfyReloadObjectInfo(false);
     }
+    if (isVisible(getRequiredElementById('Server-Info'))) {
+        comfyTorchManager.refresh();
+    }
 });
 
 function comfyListWorkflowsForBrowser(path, isRefresh, callback, depth) {

src/BuiltinExtensions/ComfyUIBackend/ComfyUISelfStartBackend.cs

@@ -87,6 +87,9 @@ public class ComfyUISelfStartSettings : AutoConfiguration
 
     public static string SwarmValidatedFrontendVersion = "1.42.11";
 
+    /// <summary>The current known version of PyTorch.</summary>
+    public static string CurrentTorchVersion = "2.12.1";
+
     /// <summary>List of known required python packages, as pairs of strings: Item1 is the folder name within python packages to look for, Item2 is the pip install command.</summary>
     public static List<(string, string)> RequiredPythonPackages =
     [
@@ -721,6 +724,26 @@ public static Version ParseVersion(string vers)
         return Version.Parse(vers.Before(".dev"));
     }
 
+    /// <summary>Get the version of a single installed pip package.</summary>
+    public static string GetInstalledPackageVersion(string startScript, string package)
+    {
+        string lib = NetworkBackendUtils.GetProbableLibFolderFor(startScript);
+        if (lib is null || lib.Length < 3 || !Directory.Exists(lib))
+        {
+            return null;
+        }
+        string prefix = $"{package}-";
+        foreach (string dir in Directory.EnumerateDirectories(lib))
+        {
+            string name = dir.Replace('\\', '/').AfterLast('/');
+            if (name.EndsWith(".dist-info") && name.StartsWith(prefix))
+            {
+                return name[prefix.Length..].Before(".dist-info");
+            }
+        }
+        return null;
+    }
+
     /// <summary>Strict matcher that will block any muckery, excluding URLs and etc.</summary>
     public static AsciiMatcher RequirementPartMatcher = new(AsciiMatcher.BothCaseLetters + AsciiMatcher.Digits + ".-_");
 

src/BuiltinExtensions/ComfyUIBackend/ComfyUIWebAPI.cs

@@ -8,6 +8,7 @@
 using SwarmUI.Text2Image;
 using SwarmUI.Utils;
 using SwarmUI.WebAPI;
+using System.Diagnostics;
 using System.IO;
 using System.Net.WebSockets;
 
@@ -26,6 +27,8 @@ public static void Register()
         API.RegisterAPICall(ComfyGetNodeTypesForBackend, false, Permissions.ViewBackendsList);
         API.RegisterAPICall(ComfyEnsureRefreshable, false, ComfyUIBackendExtension.PermDirectCalls);
         API.RegisterAPICall(ComfyInstallFeatures, true, Permissions.InstallFeatures);
+        API.RegisterAPICall(ComfyListTorchInstalls, false, Permissions.InstallFeatures);
+        API.RegisterAPICall(ComfyUpdateTorch, true, Permissions.InstallFeatures);
         API.RegisterAPICall(DoTensorRTCreateWS, true, Permissions.CreateTRT);
     }
 
@@ -243,6 +246,108 @@ public static async Task<JObject> ComfyInstallFeatures(Session session, string f
         }
     }
 
+    /// <summary>API route to list torch version installs for comfy backends.</summary>
+    public static async Task<JObject> ComfyListTorchInstalls(Session session)
+    {
+        Version target = ComfyUISelfStartBackend.ParseVersion(ComfyUISelfStartBackend.CurrentTorchVersion);
+        Dictionary<string, List<int>> folders = [];
+        foreach (BackendHandler.BackendData data in Program.Backends.EnumerateT2IBackends)
+        {
+            if (data.AbstractBackend is not ComfyUISelfStartBackend backend)
+            {
+                continue;
+            }
+            string script = backend.Settings.StartScript;
+            if (string.IsNullOrWhiteSpace(script))
+            {
+                continue;
+            }
+            string folder = backend.ComfyPathBase;
+            folders.GetOrCreate(folder, () => []).Add(data.ID);
+        }
+        JArray installs = [];
+        foreach ((string folder, List<int> backends) in folders.OrderBy(p => p.Key))
+        {
+            string torchVersRaw = ComfyUISelfStartBackend.GetInstalledPackageVersion($"{folder}/main.py", "torch");
+            string torchVers = torchVersRaw?.Before('+');
+            bool canUpdate = false;
+            if (torchVers is not null && target is not null)
+            {
+                try
+                {
+                    canUpdate = ComfyUISelfStartBackend.ParseVersion(torchVers) < target;
+                }
+                catch (Exception)
+                {
+                    canUpdate = false;
+                }
+                if (!torchVersRaw.ToLowerFast().Contains("+cu")) // TODO: AMD/etc support?
+                {
+                    canUpdate = false;
+                }
+            }
+            installs.Add(new JObject()
+            {
+                ["path"] = folder,
+                ["torch_version"] = torchVersRaw ?? "(unknown)",
+                ["can_update"] = canUpdate,
+                ["backend_ids"] = new JArray(backends.OrderBy(i => i).Select(i => (JToken)i).ToArray())
+            });
+        }
+        return new JObject()
+        {
+            ["target_torch_version"] = ComfyUISelfStartBackend.CurrentTorchVersion,
+            ["installs"] = installs
+        };
+    }
+
+    /// <summary>API route to update Torch for a single Comfy install folder.</summary>
+    public static async Task<JObject> ComfyUpdateTorch(Session session, int backendId)
+    {
+        await MultiInstallLock.WaitAsync(Program.GlobalProgramCancel);
+        try
+        {
+            if (!Program.Backends.AllBackends.TryGetValue(backendId, out BackendHandler.BackendData backend) || backend.AbstractBackend is not ComfyUISelfStartBackend target)
+            {
+                return new() { ["error"] = $"Invalid backend ID {backendId}" };
+            }
+            string path = target.ComfyPathBase;
+            ComfyUISelfStartBackend[] backends = [.. Program.Backends.EnumerateT2IBackends.Select(d => d.AbstractBackend as ComfyUISelfStartBackend).Where(b => b is not null && b.ComfyPathBase == path)];
+            Logs.Info($"User {session.User.UserID} requested a PyTorch update for ComfyUI install at '{path}' (affecting {backends.Length} backend(s))...");
+            Logs.Info($"[Torch Update] Shutting down {backends.Length} backend(s) that use this install...");
+            Task[] shutdownTasks = [.. backends.Select(b => Program.Backends.ShutdownBackendCleanly(b.BackendData))];
+            await Task.WhenAll(shutdownTasks);
+            async Task pipCall(string reason, string call)
+            {
+                Logs.Info($"[Torch Update] {reason} for ComfyUI install '{path}'...");
+                Process p = target.DoPythonCall($"-s -m pip {call}");
+                NetworkBackendUtils.ReportLogsFromProcess(p, $"ComfyUI (Torch Update - {reason})", "");
+                await p.WaitForExitAsync(Program.GlobalProgramCancel);
+                Logs.Info($"[Torch Update] Done {reason} for ComfyUI install '{path}'.");
+            }
+            try
+            {
+                await pipCall("Uninstalling old torch", "uninstall -y torch torchvision torchaudio");
+                await pipCall("Installing new torch", "install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu132");
+            }
+            catch (Exception ex)
+            {
+                Logs.Error($"Failed to update torch for ComfyUI install '{path}': {ex.ReadableString()}");
+                return new JObject() { ["error"] = "Torch update failed, check server logs for details." };
+            }
+            Logs.Info($"[Torch Update] Restarting {backends.Length} backend(s)...");
+            foreach (ComfyUISelfStartBackend back in backends)
+            {
+                Program.Backends.DoInitBackend(back.BackendData);
+            }
+            return new() { ["success"] = true };
+        }
+        finally
+        {
+            MultiInstallLock.Release();
+        }
+    }
+
     public static Dictionary<string, int> AspectRangeToMultiplier = new()
     {
         ["Exact"] = 1,