Merge pull request LykosAI#984 from ionite34/blackwell-support

Add blackwell torch/vision versions for comfyUI & fix unet path for l…

(cherry picked from commit ef90eb2)

# Conflicts:
#	CHANGELOG.md

Author: mohnjiles

CHANGELOG.md

@@ -5,6 +5,71 @@ 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).
 
+<<<<<<< HEAD
+=======
+## v2.14.0-dev.3
+### 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.14.0-dev.2
+### Added
+- Added Align Your Steps scheduler to Inference
+- Added wildcards to Inference prompts, e.g. `{blue|green|red}` will randomly select one of the colors
+- Added Safetensor Metadata viewer to the Checkpoint Manager context menu - thanks to @genteure!
+### Changed
+- Updated the Civitai Model Browser base model selector to match the new Checkpoint Manager filter UI 
+- FaceDetailers in Inference will now inherit the primary sampler/scheduler/etc. by default. You can still manually set these by enabling the options via the ⚙️ button on the FaceDetailer card
+- Slightly rearranged the FaceDetailer card layout due to the above change
+- "Remove symbolic links on shutdown" option now also removes links from Output Sharing
+- Inference "Extra Networks" selector now filters extra networks based on the selected base model
+- Updated Japanese, Brazilian Portuguese, Chinese, and Russian translations
+### Fixed
+- Fixed crash when dragging & dropping images in Inference (hopefully)
+- Fixed HiresFix Inference addon not inheriting sampler/scheduler properly
+- Fixed some plus (+) buttons getting cut off in the Inference UI
+- Fixed CFG Rescale addon interfering with refiner model in Inference
+- Fixed [#1083](https://github.com/LykosAI/StabilityMatrix/issues/1083) - "Show Nested Models" incorrectly displaying models from some non-nested folders
+- Fixed issue with InvokeAI model sharing when the host address is set to 0.0.0.0
+- Fixed issue when parsing index URLs in Python Dependencies Override menu
+- Fixed ComfyUI-Zluda not respecting pip user overrides
+- Fixed issue with Checkpoint Manager not displaying any models
+- (dev.2 re-release) Fixed autocomplete not showing in certain cases when using wildcards
+- (dev.2 re-release) Fixed package restart button not working 
+- (dev.2 re-release) Fixed [#1120](https://github.com/LykosAI/StabilityMatrix/issues/1120) - crash when right clicking in the console after restarting a package
+### Supporters
+#### Visionaries
+- A huge thank you to our incredible Visionary-tier Patreon supporters, **Waterclouds**, **TheTekknician**, and our newest Visionary, **Corey**! Your generous support is greatly appreciated!
+
+## v2.14.0-dev.1
+### Added
+- Added Rescale CFG addon to Inference
+- Added Swap Dimensions button between the width/height input in Inference
+- Added Ctrl+Tab/Ctrl+Shift+Tab shortcuts for navigating between Inference tabs
+- Added OpenModelDB tab to the Model Browser
+### Changed
+- Improved the quality of Inference inpainting by upgrading the workflow behind the scenes. The workflow remains the same for you — just better results!
+- Redesigned the Checkpoint Manager Filter flyout to include more options and improve the layout
+- "Clear All" button will now remain at the top of the Downloads list regardless of scroll position - thanks to @Genteure!
+- Improved image metadata parsing - thanks to @Genteure!
+### Fixed
+- Fixed Inference image selector card buttons taking up the whole height of the card
+- Fixed Inference mask editor failing to paint to the right-most edge on large images
+- Fixed Inference mask editor not showing the entire image in certain circumstances
+- Fixed an issue where certain sampler/scheduler combos would not get saved in image metadata - thanks to @yansigit!
+- Fixed [#1078](https://github.com/LykosAI/StabilityMatrix/issues/1078) - "Call from invalid thread" error after one-click install finishes
+- Fixed [#1080](https://github.com/LykosAI/StabilityMatrix/issues/1080) - Some models not displayed in Checkpoint Manager
+### Supporters
+#### Visionaries
+- Many thanks to our incredible Visionary-tier Patreon supporters, **Waterclouds** and **TheTekknician**! Your support helps us continue to improve Stability Matrix!
+
+>>>>>>> ef90eb2f (Merge pull request #984 from ionite34/blackwell-support)
 ## v2.13.4
 ### Changed
 - Upgraded ComfyUI CUDA torch to 12.6

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") }
                 }
             );
         }