Merge pull request #437 from ionite34/add-onetrainer

add onetrainer & moved tk/tcl env vars to BaseGitPackage

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 
 ## v2.8.0-dev.3
 ### Added
+- Added new package: [OneTrainer](https://github.com/Nerogar/OneTrainer)
 - Added ComfyUI launch argument configs: Cross Attention Method, Force Floating Point Precision, VAE Precision
 - Added Delete button to the CivitAI Model Browser details dialog
 - Added "Copy Link to Clipboard" for connected models in the Checkpoints page

StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs

@@ -2,6 +2,7 @@
 using System.IO.Compression;
 using NLog;
 using Octokit;
+using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Helper.Cache;
 using StabilityMatrix.Core.Models.Database;
@@ -147,10 +148,24 @@ public abstract class BaseGitPackage : BasePackage
             await VenvRunner.DisposeAsync().ConfigureAwait(false);
         }
 
+        // Set additional required environment variables
+        var env = new Dictionary<string, string>();
+        if (SettingsManager.Settings.EnvironmentVariables is not null)
+        {
+            env.Update(SettingsManager.Settings.EnvironmentVariables);
+        }
+
+        if (Compat.IsWindows)
+        {
+            var tkPath = Path.Combine(SettingsManager.LibraryDir, "Assets", "Python310", "tcl", "tcl8.6");
+            env["TCL_LIBRARY"] = tkPath;
+            env["TK_LIBRARY"] = tkPath;
+        }
+
         VenvRunner = new PyVenvRunner(venvPath)
         {
             WorkingDirectory = installedPackagePath,
-            EnvironmentVariables = SettingsManager.Settings.EnvironmentVariables,
+            EnvironmentVariables = env
         };
 
         if (!VenvRunner.Exists() || forceRecreate)

StabilityMatrix.Core/Models/Packages/KohyaSs.cs

@@ -40,14 +40,16 @@ public class KohyaSs(
 
     public override TorchVersion GetRecommendedTorchVersion() => TorchVersion.Cuda;
 
-    public override string Disclaimer => "Nvidia GPU with at least 8GB VRAM is recommended. May be unstable on Linux.";
+    public override string Disclaimer =>
+        "Nvidia GPU with at least 8GB VRAM is recommended. May be unstable on Linux.";
 
     public override PackageDifficulty InstallerSortOrder => PackageDifficulty.UltraNightmare;
 
     public override bool OfferInOneClickInstaller => false;
     public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None;
-    public override IEnumerable<TorchVersion> AvailableTorchVersions => new[] { TorchVersion.Cuda };
-    public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods => new[] { SharedFolderMethod.None };
+    public override IEnumerable<TorchVersion> AvailableTorchVersions => [TorchVersion.Cuda];
+    public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods =>
+        new[] { SharedFolderMethod.None };
 
     public override List<LaunchOptionDefinition> LaunchOptions =>
         [
@@ -189,7 +191,9 @@ public class KohyaSs(
                 """
             );
 
-            var replacementAcceleratePath = Compat.IsWindows ? @".\venv\scripts\accelerate" : "./venv/bin/accelerate";
+            var replacementAcceleratePath = Compat.IsWindows
+                ? @".\venv\scripts\accelerate"
+                : "./venv/bin/accelerate";
 
             var replacer = scope.InvokeMethod(
                 "StringReplacer",
@@ -238,7 +242,6 @@ public class KohyaSs(
 
         var args = $"\"{Path.Combine(installedPackagePath, command)}\" {arguments}";
 
-        VenvRunner.EnvironmentVariables = GetEnvVars(installedPackagePath);
         VenvRunner.RunDetached(args.TrimEnd(), HandleConsoleOutput, OnExit);
     }
 
@@ -246,23 +249,4 @@ public class KohyaSs(
     public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders { get; }
 
     public override string MainBranch => "master";
-
-    private Dictionary<string, string> GetEnvVars(string installDirectory)
-    {
-        // Set additional required environment variables
-        var env = new Dictionary<string, string>();
-        if (SettingsManager.Settings.EnvironmentVariables is not null)
-        {
-            env.Update(SettingsManager.Settings.EnvironmentVariables);
-        }
-
-        if (!Compat.IsWindows)
-            return env;
-
-        var tkPath = Path.Combine(SettingsManager.LibraryDir, "Assets", "Python310", "tcl", "tcl8.6");
-        env["TCL_LIBRARY"] = tkPath;
-        env["TK_LIBRARY"] = tkPath;
-
-        return env;
-    }
 }

StabilityMatrix.Core/Models/Packages/OneTrainer.cs

@@ -0,0 +1,92 @@
+using System.Diagnostics;
+using System.Text.RegularExpressions;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Helper.Cache;
+using StabilityMatrix.Core.Models.Progress;
+using StabilityMatrix.Core.Processes;
+using StabilityMatrix.Core.Python;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Core.Models.Packages;
+
+[Singleton(typeof(BasePackage))]
+public class OneTrainer(
+    IGithubApiCache githubApi,
+    ISettingsManager settingsManager,
+    IDownloadService downloadService,
+    IPrerequisiteHelper prerequisiteHelper
+) : BaseGitPackage(githubApi, settingsManager, downloadService, prerequisiteHelper)
+{
+    public override string Name => "OneTrainer";
+    public override string DisplayName { get; set; } = "OneTrainer";
+    public override string Author => "Nerogar";
+    public override string Blurb =>
+        "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 Uri PreviewImageUri =>
+        new("https://github.com/Nerogar/OneTrainer/blob/master/resources/icons/icon.png?raw=true");
+
+    public override string OutputFolderName => string.Empty;
+    public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None;
+    public override IEnumerable<TorchVersion> AvailableTorchVersions => [TorchVersion.Cuda];
+    public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods =>
+        new[] { SharedFolderMethod.None };
+    public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Nightmare;
+    public override bool OfferInOneClickInstaller => false;
+    public override bool ShouldIgnoreReleases => true;
+
+    public override async Task InstallPackage(
+        string installLocation,
+        TorchVersion torchVersion,
+        SharedFolderMethod selectedSharedFolderMethod,
+        DownloadPackageVersionOptions versionOptions,
+        IProgress<ProgressReport>? progress = null,
+        Action<ProcessOutput>? onConsoleOutput = null
+    )
+    {
+        progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true));
+
+        await using var venvRunner = new PyVenvRunner(Path.Combine(installLocation, "venv"));
+        venvRunner.WorkingDirectory = installLocation;
+        await venvRunner.Setup(true, onConsoleOutput).ConfigureAwait(false);
+
+        progress?.Report(new ProgressReport(-1f, "Installing requirements", isIndeterminate: true));
+
+        var pipArgs = new PipInstallArgs("-r", "requirements.txt");
+        await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false);
+    }
+
+    public override async Task RunPackage(
+        string installedPackagePath,
+        string command,
+        string arguments,
+        Action<ProcessOutput>? onConsoleOutput
+    )
+    {
+        await SetupVenv(installedPackagePath).ConfigureAwait(false);
+        var args = $"\"{Path.Combine(installedPackagePath, command)}\" {arguments}";
+
+        VenvRunner?.RunDetached(args.TrimEnd(), HandleConsoleOutput, HandleExit);
+        return;
+
+        void HandleExit(int i)
+        {
+            Debug.WriteLine($"Venv process exited with code {i}");
+            OnExit(i);
+        }
+
+        void HandleConsoleOutput(ProcessOutput s)
+        {
+            onConsoleOutput?.Invoke(s);
+        }
+    }
+
+    public override List<LaunchOptionDefinition> LaunchOptions => [LaunchOptionDefinition.Extras];
+    public override Dictionary<SharedFolderType, IReadOnlyList<string>>? SharedFolders { get; }
+    public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders { get; }
+    public override string MainBranch => "master";
+}