Merge pull request #1684 from NeuralFault/win-rocm-onetrainer

feat: Integrate OneTrainer into Windows ROCm Helper [Windows/AMD]

Author: mohnjiles

StabilityMatrix.Core/Helper/Factory/PackageFactory.cs

@@ -120,7 +120,8 @@ public BasePackage GetNewBasePackage(InstalledPackage installedPackage)
                 downloadService,
                 prerequisiteHelper,
                 pyInstallationManager,
-                pipWheelService
+                pipWheelService,
+                rocmPackageHelper
             ),
             "RuinedFooocus" => new RuinedFooocus(
                 githubApiCache,

StabilityMatrix.Core/Models/Packages/OneTrainer.cs

@@ -5,9 +5,11 @@
 using StabilityMatrix.Core.Helper.HardwareInfo;
 using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Models.Progress;
+using StabilityMatrix.Core.Models.Rocm;
 using StabilityMatrix.Core.Processes;
 using StabilityMatrix.Core.Python;
 using StabilityMatrix.Core.Services;
+using StabilityMatrix.Core.Services.Rocm;
 
 namespace StabilityMatrix.Core.Models.Packages;
 
@@ -18,7 +20,8 @@ public class OneTrainer(
     IDownloadService downloadService,
     IPrerequisiteHelper prerequisiteHelper,
     IPyInstallationManager pyInstallationManager,
-    IPipWheelService pipWheelService
+    IPipWheelService pipWheelService,
+    IRocmPackageHelper rocmPackageHelper
 )
     : BaseGitPackage(
         githubApi,
@@ -36,7 +39,7 @@ IPipWheelService pipWheelService
         "OneTrainer is a one-stop solution for all your stable diffusion training needs";
     public override string LicenseType => "AGPL-3.0";
     public override string LicenseUrl => "https://github.com/Nerogar/OneTrainer/blob/master/LICENSE.txt";
-    public override string LaunchCommand => "scripts/train_ui.py";
+    public override string LaunchCommand => "scripts/train_ui_ctk.py";
 
     public override Uri PreviewImageUri =>
         new("https://github.com/Nerogar/OneTrainer/blob/master/resources/icons/icon.png?raw=true");
@@ -52,6 +55,17 @@ IPipWheelService pipWheelService
     public override bool ShouldIgnoreReleases => true;
     public override IEnumerable<PackagePrerequisite> Prerequisites =>
         base.Prerequisites.Concat([PackagePrerequisite.Tkinter]);
+    public override PyVersion RecommendedPythonVersion => Python.PyInstallationManager.Python_3_12_10;
+
+    public override TorchIndex GetRecommendedTorchVersion()
+    {
+        if (Compat.IsWindows && rocmPackageHelper.GetCompatibility().IsCompatible)
+        {
+            return TorchIndex.Rocm;
+        }
+
+        return base.GetRecommendedTorchVersion();
+    }
 
     public override async Task InstallPackage(
         string installLocation,
@@ -69,17 +83,66 @@ public override async Task InstallPackage(
             )
             .ConfigureAwait(false);
 
-        progress?.Report(new ProgressReport(-1f, "Installing requirements", isIndeterminate: true));
         var torchVersion = options.PythonOptions.TorchIndex ?? GetRecommendedTorchVersion();
-        var requirementsFileName = torchVersion switch
+        var pyVersion = options.PythonOptions.PythonVersion ?? RecommendedPythonVersion;
+
+        // Windows ROCm path
+        var isWindowsRocm =
+            Compat.IsWindows
+            && torchVersion == TorchIndex.Rocm
+            && rocmPackageHelper.GetCompatibility().IsCompatible;
+
+        if (isWindowsRocm)
+        {
+            var config = rocmPackageHelper.BuildWindowsNativeInstallConfig(
+                OneTrainerWindowsRocmProfile.Default
+            );
+
+            await StandardPipInstallProcessAsync(
+                    venvRunner,
+                    options,
+                    installedPackage,
+                    config with
+                    {
+                        RequirementsFilePaths = ["requirements-default.txt"],
+                    },
+                    onConsoleOutput,
+                    progress,
+                    cancellationToken
+                )
+                .ConfigureAwait(false);
+
+            progress?.Report(
+                new ProgressReport(-1f, "Installing Windows ROCm torch...", isIndeterminate: true)
+            );
+
+            await rocmPackageHelper
+                .InstallWindowsNativeTorchAsync(
+                    venvRunner,
+                    installedPackage,
+                    OneTrainerWindowsRocmProfile.CreateInstallProfile(pyVersion),
+                    progress,
+                    onConsoleOutput,
+                    cancellationToken
+                )
+                .ConfigureAwait(false);
+        }
+        else
         {
-            TorchIndex.Cuda => "requirements-cuda.txt",
-            TorchIndex.Rocm => "requirements-rocm.txt",
-            _ => "requirements-default.txt",
-        };
+            // Original path (CUDA / Linux ROCm)
+            progress?.Report(new ProgressReport(-1f, "Installing requirements", isIndeterminate: true));
 
-        await venvRunner.PipInstall(["-r", requirementsFileName], onConsoleOutput).ConfigureAwait(false);
+            var requirementsFileName = torchVersion switch
+            {
+                TorchIndex.Cuda => "requirements-cuda.txt",
+                TorchIndex.Rocm => "requirements-rocm.txt",
+                _ => "requirements-default.txt",
+            };
 
+            await venvRunner.PipInstall(["-r", requirementsFileName], onConsoleOutput).ConfigureAwait(false);
+        }
+
+        // Shared requirements-global.txt (both paths)
         var requirementsGlobal = new FilePath(installLocation, "requirements-global.txt");
         var pipArgs = new PipInstallArgs().WithParsedFromRequirementsTxt(
             (await requirementsGlobal.ReadAllTextAsync(cancellationToken).ConfigureAwait(false)).Replace(
@@ -108,13 +171,33 @@ public override async Task RunPackage(
         await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
             .ConfigureAwait(false);
 
+        var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion();
+
+        if (rocmPackageHelper.ShouldApplyWindowsLaunchEnvironment(selectedTorchIndex))
+        {
+            VenvRunner.UpdateEnvironmentVariables(env =>
+                env.SetItems(rocmPackageHelper.BuildLaunchEnvironment(OneTrainerWindowsRocmProfile.Default))
+            );
+        }
+
+        foreach (var line in GetLaunchNoticeLines(installedPackage))
+        {
+            onConsoleOutput?.Invoke(ProcessOutput.FromStdOutLine($"{line}{Environment.NewLine}"));
+        }
+
         VenvRunner.RunDetached(
             [Path.Combine(installLocation, options.Command ?? LaunchCommand), .. options.Arguments],
             onConsoleOutput,
             OnExit
         );
     }
 
+    private IReadOnlyList<string> GetLaunchNoticeLines(InstalledPackage installedPackage)
+    {
+        var selectedTorchIndex = installedPackage.PreferredTorchIndex ?? GetRecommendedTorchVersion();
+        return rocmPackageHelper.GetWindowsLaunchNoticeLines(selectedTorchIndex);
+    }
+
     public override List<LaunchOptionDefinition> LaunchOptions => [LaunchOptionDefinition.Extras];
     public override Dictionary<SharedFolderType, IReadOnlyList<string>>? SharedFolders { get; }
     public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders { get; }

StabilityMatrix.Core/Models/Rocm/OneTrainerWindowsRocmProfile.cs

@@ -0,0 +1,39 @@
+using StabilityMatrix.Core.Models.Packages;
+using StabilityMatrix.Core.Python;
+
+namespace StabilityMatrix.Core.Models.Rocm;
+
+/// <summary>
+/// Shared Windows ROCm profile for OneTrainer.
+/// Python 3.12 only - required by compatible bitsandbytes wheel.
+/// </summary>
+public class OneTrainerWindowsRocmProfile : RocmPackageProfile
+{
+    public static RocmPackageProfile Default { get; } = new OneTrainerWindowsRocmProfile();
+
+    // Restores flop counter functionality requiring triton module
+    private const string TritonWindowsPackage = "triton-windows";
+
+    // Replace upstream bitsandbytes with ROCm-aware bitsandbytes for ROCm Technical Preview on Windows
+    private const string BitsAndBytesWheelUrl =
+        "https://github.com/0xDELUXA/bitsandbytes_win_rocm/releases/download/0.50.0.dev0-py3.12-rocm7.15-win_amd64_all/bitsandbytes-0.50.0.dev0-cp312-cp312-win_amd64.whl";
+
+    public static RocmPackageProfile CreateInstallProfile(PyVersion pyVersion)
+    {
+        if (pyVersion.Major == 3 && pyVersion.Minor == 12)
+        {
+            return new RocmPackageProfile
+            {
+                InstallConfig = new PipInstallConfig
+                {
+                    PostTorchInstallPipArgs = [TritonWindowsPackage, BitsAndBytesWheelUrl],
+                },
+            };
+        }
+
+        return new RocmPackageProfile
+        {
+            InstallConfig = new PipInstallConfig { PostTorchInstallPipArgs = [TritonWindowsPackage] },
+        };
+    }
+}