Merge pull request LykosAI#985 from ionite34/backport/main/pr-984

[dev to main] backport: Add blackwell torch/vision versions for comfyUI & fix unet path for l… (984)

Author: mohnjiles

CHANGELOG.md

@@ -4,14 +4,17 @@ All notable changes to Stability Matrix will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2.0.0.html).
-
+  
 ## v2.13.4
+### Added
+- Added support for RTX 5000-series GPUs in ComfyUI
 ### Changed
 - Upgraded ComfyUI CUDA torch to 12.6
 ### Fixed
 - Fixed [#1128](https://github.com/LykosAI/StabilityMatrix/issues/1128) - overwriting models when downloading multiple with the same name
 - Fixed ROCm torch indexes for ComfyUI & Forge
 - Fixed model browser sometimes downloading to `ModelsLora` or `ModelsStableDiffusion` folders instead of the correct folder
+- Fixed incorrect Unet folder path for ComfyUI users on Linux/macOS
 
 ## v2.13.3
 ### Changed

StabilityMatrix.Core/Helper/HardwareInfo/GpuInfo.cs

@@ -27,6 +27,16 @@ public bool IsNvidia
         }
     }
 
+    public bool IsBlackwellGpu()
+    {
+        if (Name is null)
+            return false;
+
+        return IsNvidia
+            && Name.Contains("RTX 50", StringComparison.OrdinalIgnoreCase)
+            && !Name.Contains("RTX 5000", StringComparison.OrdinalIgnoreCase);
+    }
+
     public bool IsAmd => Name?.Contains("amd", StringComparison.OrdinalIgnoreCase) ?? false;
     public bool IsIntel => Name?.Contains("arc", StringComparison.OrdinalIgnoreCase) ?? false;
 

StabilityMatrix.Core/Helper/HardwareInfo/HardwareHelper.cs

@@ -250,6 +250,17 @@ public static bool HasNvidiaGpu()
         return IterGpuInfo().Any(gpu => gpu.IsNvidia);
     }
 
+    public static bool HasBlackwellGpu()
+    {
+        return IterGpuInfo()
+            .Any(
+                gpu =>
+                    gpu is { IsNvidia: true, Name: not null }
+                    && gpu.Name.Contains("RTX 50", StringComparison.OrdinalIgnoreCase)
+                    && !gpu.Name.Contains("RTX 5000", StringComparison.OrdinalIgnoreCase)
+            );
+    }
+
     /// <summary>
     /// Return true if the system has at least one AMD GPU.
     /// </summary>

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -207,36 +207,52 @@ public override async Task InstallPackage(
 
         var pipArgs = new PipInstallArgs();
 
-        pipArgs = torchVersion switch
+        var isBlackwell =
+            SettingsManager.Settings.PreferredGpu?.IsBlackwellGpu() ?? HardwareHelper.HasBlackwellGpu();
+
+        if (isBlackwell)
         {
-            TorchIndex.DirectMl => pipArgs.WithTorchDirectML(),
-            _
-                => pipArgs
-                    .AddArg("--upgrade")
-                    .WithTorch()
-                    .WithTorchVision()
-                    .WithTorchExtraIndex(
-                        torchVersion switch
-                        {
-                            TorchIndex.Cpu => "cpu",
-                            TorchIndex.Cuda => "cu126",
-                            TorchIndex.Rocm => "rocm6.2.4",
-                            TorchIndex.Mps => "cpu",
-                            _
-                                => throw new ArgumentOutOfRangeException(
-                                    nameof(torchVersion),
-                                    torchVersion,
-                                    null
-                                )
-                        }
-                    )
-        };
+            pipArgs = pipArgs
+                .AddArg(
+                    "https://huggingface.co/w-e-w/torch-2.6.0-cu128.nv/resolve/main/torch-2.6.0%2Bcu128.nv-cp310-cp310-win_amd64.whl"
+                )
+                .AddArg(
+                    "https://huggingface.co/w-e-w/torch-2.6.0-cu128.nv/resolve/main/torchvision-0.20.0a0%2Bcu128.nv-cp310-cp310-win_amd64.whl"
+                );
+        }
+        else
+        {
+            pipArgs = torchVersion switch
+            {
+                TorchIndex.DirectMl => pipArgs.WithTorchDirectML(),
+                _
+                    => pipArgs
+                        .AddArg("--upgrade")
+                        .WithTorch()
+                        .WithTorchVision()
+                        .WithTorchExtraIndex(
+                            torchVersion switch
+                            {
+                                TorchIndex.Cpu => "cpu",
+                                TorchIndex.Cuda => "cu126",
+                                TorchIndex.Rocm => "rocm6.2.4",
+                                TorchIndex.Mps => "cpu",
+                                _
+                                    => throw new ArgumentOutOfRangeException(
+                                        nameof(torchVersion),
+                                        torchVersion,
+                                        null
+                                    )
+                            }
+                        )
+            };
+        }
 
         var requirements = new FilePath(installLocation, "requirements.txt");
 
         pipArgs = pipArgs.WithParsedFromRequirementsTxt(
             await requirements.ReadAllTextAsync(cancellationToken).ConfigureAwait(false),
-            excludePattern: "torch$|numpy"
+            excludePattern: isBlackwell ? "torch$|torchvision$|numpy" : "torch$|numpy"
         );
 
         // https://github.com/comfyanonymous/ComfyUI/pull/4121
@@ -385,7 +401,7 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
             nodeValue.Children["ultralytics_bbox"] = Path.Combine(modelsDir, "Ultralytics", "bbox");
             nodeValue.Children["ultralytics_segm"] = Path.Combine(modelsDir, "Ultralytics", "segm");
             nodeValue.Children["sams"] = Path.Combine(modelsDir, "Sams");
-            nodeValue.Children["diffusion_models"] = Path.Combine(modelsDir, "unet");
+            nodeValue.Children["diffusion_models"] = Path.Combine(modelsDir, "Unet");
         }
         else
         {
@@ -429,7 +445,7 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
                     { "ultralytics_bbox", Path.Combine(modelsDir, "Ultralytics", "bbox") },
                     { "ultralytics_segm", Path.Combine(modelsDir, "Ultralytics", "segm") },
                     { "sams", Path.Combine(modelsDir, "Sams") },
-                    { "diffusion_models", Path.Combine(modelsDir, "unet") }
+                    { "diffusion_models", Path.Combine(modelsDir, "Unet") }
                 }
             );
         }