Browse Source

Merge pull request #437 from ionite34/add-onetrainer

add onetrainer & moved tk/tcl env vars to BaseGitPackage
pull/438/head
JT 11 months ago committed by GitHub
parent
commit
9b878d4ef3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 1
      CHANGELOG.md
  2. 17
      StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs
  3. 32
      StabilityMatrix.Core/Models/Packages/KohyaSs.cs
  4. 92
      StabilityMatrix.Core/Models/Packages/OneTrainer.cs

1
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

17
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)

32
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;
}
}

92
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";
}
Loading…
Cancel
Save