Add Fizgig LoRA training studio package #1761
+208
−2
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,194 @@ | ||
| using Injectio.Attributes; | ||
|
Member
There was a problem hiding this comment. nit: tiny one, the file starts with a UTF-8 BOM. A few older package files have one too, but we're avoiding it in new files. |
||
| using StabilityMatrix.Core.Helper; | ||
| using StabilityMatrix.Core.Helper.Cache; | ||
| using StabilityMatrix.Core.Helper.HardwareInfo; | ||
| using StabilityMatrix.Core.Models.Progress; | ||
| using StabilityMatrix.Core.Processes; | ||
| using StabilityMatrix.Core.Python; | ||
| using StabilityMatrix.Core.Services; | ||
|
|
||
| namespace StabilityMatrix.Core.Models.Packages; | ||
|
|
||
| [RegisterSingleton<BasePackage, Fizgig>(Duplicate = DuplicateStrategy.Append)] | ||
| public class Fizgig( | ||
| IGithubApiCache githubApi, | ||
| ISettingsManager settingsManager, | ||
| IDownloadService downloadService, | ||
| IPrerequisiteHelper prerequisiteHelper, | ||
| IPyInstallationManager pyInstallationManager, | ||
| IPipWheelService pipWheelService | ||
| ) | ||
| : BaseGitPackage( | ||
| githubApi, | ||
| settingsManager, | ||
| downloadService, | ||
| prerequisiteHelper, | ||
| pyInstallationManager, | ||
| pipWheelService | ||
| ) | ||
| { | ||
| public override string Name => "Fizgig"; | ||
| public override string DisplayName { get; set; } = "Fizgig"; | ||
| public override string Author => "shootthesound"; | ||
|
|
||
| public override string Blurb => | ||
| "LoRA training studio for Flux 2 Klein 9B, Krea 2, MiniMax H3 and Qwen Image 2.1 — train, profile, repair and extract"; | ||
|
Member
There was a problem hiding this comment. nit: we keep em dashes out of user-facing copy. Maybe |
||
|
|
||
| // Shown in the install browser before the user commits to installing. | ||
| public override string Disclaimer => | ||
| Compat.IsWindows | ||
| ? "Visual Studio Build Tools for C++ Desktop Development will be installed system-wide if not already present (may require admin privileges). " | ||
| + "They are shared with other software and remain installed after Fizgig is uninstalled." | ||
| : string.Empty; | ||
|
|
||
| public override string LicenseType => "Apache-2.0"; | ||
| public override string LicenseUrl => "https://github.com/shootthesound/Fizgig/blob/master/LICENSE"; | ||
|
|
||
| // NOT launch.pyw: that launcher re-spawns itself under venv/Scripts/pythonw.exe and exits, | ||
| // which would drop the process we track (no console output, no working Stop button) and | ||
| // leave the GUI orphaned. lora_trainer_gui.py has a standalone main() and is what upstream's | ||
| // run_fizgig.sh invokes directly. | ||
| public override string LaunchCommand => "lora_trainer_gui.py"; | ||
|
|
||
| public override Uri PreviewImageUri => | ||
| new("https://github.com/shootthesound/Fizgig/blob/master/icon.png?raw=true"); | ||
|
|
||
| public override string MainBranch => "master"; | ||
| public override PackageType PackageType => PackageType.SdTraining; | ||
| public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Advanced; | ||
| public override bool OfferInOneClickInstaller => false; | ||
| public override bool IsCompatible => HardwareHelper.HasNvidiaGpu(); | ||
| public override IEnumerable<TorchIndex> AvailableTorchIndices => [TorchIndex.Cuda]; | ||
|
|
||
| public override TorchIndex GetRecommendedTorchVersion() => TorchIndex.Cuda; | ||
|
|
||
| public override PyVersion RecommendedPythonVersion => Python.PyInstallationManager.Python_3_12_10; | ||
|
|
||
| // Tkinter for the GUI itself; VcBuildTools for triton / torch.compile's inductor backend, | ||
| // which the Compile Blocks speedup needs on Windows. | ||
| public override IEnumerable<PackagePrerequisite> Prerequisites => | ||
| base.Prerequisites.Concat([PackagePrerequisite.Tkinter, PackagePrerequisite.VcBuildTools]); | ||
|
|
||
| public override List<LaunchOptionDefinition> LaunchOptions => [LaunchOptionDefinition.Extras]; | ||
|
|
||
| // Trained LoRAs, not images. | ||
| public override string OutputFolderName => string.Empty; | ||
| public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders => null; | ||
|
|
||
| /// <summary> | ||
| /// Defaults to None, matching the other trainers. Opting in to Symlink (Package Manager -> | ||
| /// ... -> Shared Model Strategy) junctions output_loras into the shared Lora folder, so a | ||
| /// freshly trained LoRA is immediately visible to ComfyUI and friends. | ||
| /// </summary> | ||
| public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None; | ||
|
|
||
| public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods => | ||
| [SharedFolderMethod.None, SharedFolderMethod.Symlink]; | ||
|
|
||
| /// <remarks> | ||
| /// Only output_loras is mapped. Fizgig's models/ directory deliberately flattens every | ||
| /// weight it downloads into one folder — DiTs, text encoders, VAEs, turbo LoRAs and training | ||
| /// adapters all land there as bare filenames (see src/fizgig/scripts/fetch_models.py) — and | ||
| /// junctions are directory-level, so there is no way to fan that single directory out to | ||
| /// DiffusionModels/TextEncoders/VAE without them colliding on the same target path. | ||
| /// </remarks> | ||
| public override SharedFolderLayout SharedFolderLayout => | ||
| new() | ||
| { | ||
| Rules = | ||
| [ | ||
| new SharedFolderLayoutRule | ||
| { | ||
| SourceTypes = [SharedFolderType.Lora], | ||
| TargetRelativePaths = ["output_loras"], | ||
|
Member
There was a problem hiding this comment. Heads up: Fizgig writes more than finished LoRAs into With Symlink enabled, all of that ends up in the shared I'd lean towards dropping Symlink for now (None only, like the other trainers). Fizgig already has an Output Directory field, so anyone who wants their LoRAs in the shared folder can point it there themselves. |
||
| }, | ||
| ], | ||
| }; | ||
|
|
||
| public override async Task InstallPackage( | ||
| string installLocation, | ||
| InstalledPackage installedPackage, | ||
| InstallPackageOptions options, | ||
| IProgress<ProgressReport>? progress = null, | ||
| Action<ProcessOutput>? onConsoleOutput = null, | ||
| CancellationToken cancellationToken = default | ||
| ) | ||
| { | ||
| progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true)); | ||
|
|
||
| await using var venvRunner = await SetupVenvPure( | ||
| installLocation, | ||
| pythonVersion: options.PythonOptions.PythonVersion | ||
| ) | ||
| .ConfigureAwait(false); | ||
|
|
||
| // hqq ships as an sdist whose setup.py kicks off a CUDA kernel build during egg_info | ||
| // unless DISABLE_CUDA is set. Fizgig only uses its pure-PyTorch path, and its own | ||
| // requirements.txt warns never to install that line without this. | ||
| venvRunner.UpdateEnvironmentVariables(env => env.SetItem("DISABLE_CUDA", "1")); | ||
|
|
||
| const string torchVersion = "==2.10.0"; | ||
|
Member
There was a problem hiding this comment. Since releases are tracked and upstream ships often, these hardcoded pins worry me a bit. The exclude pattern strips upstream's own The risky part is triton: upstream pins Could we read the |
||
| const string torchvisionVersion = "==0.25.0"; | ||
|
|
||
| var config = new PipInstallConfig | ||
| { | ||
| RequirementsFilePaths = ["requirements.txt"], | ||
| // Drop the torch pins and the cu128 index line from the file so the pre-install step | ||
| // below is the single source of truth for which build lands in the venv. The pattern | ||
| // is anchored against the whole entry by the caller, so the version specifier has to | ||
| // be matched too - the default pattern only catches bare, unpinned names. | ||
| RequirementsExcludePattern = | ||
| @"(--extra-index-url.*|(torch|torchvision|torchaudio|xformers)([=<>!~].*)?)", | ||
| // Install the cu128 build before the requirements. accelerate (and friends) depend on | ||
| // torch transitively, so installing it afterwards would let the requirements step | ||
| // pull a default PyPI build that then has to be force-reinstalled over. | ||
| // torch 2.10 pairs with cu128 here; SM's default cu130 has no matching wheels. | ||
| PrePipInstallArgs = | ||
| [ | ||
| $"torch{torchVersion}", | ||
| $"torchvision{torchvisionVersion}", | ||
| "--extra-index-url", | ||
| "https://download.pytorch.org/whl/cu128", | ||
| "--force-reinstall", | ||
|
Member
There was a problem hiding this comment. nit: on a fresh install this step runs first, so there's nothing to force over. The flag only kicks in on updates, where it reinstalls torch, torchvision and all the |
||
| ], | ||
| // Re-state the pins alongside the requirements: pip then treats the installed | ||
| // 2.10.0+cu128 as satisfying them instead of resolving its own torch from PyPI, and | ||
| // fails loudly rather than swapping it if anything conflicts. | ||
| ExtraPipArgs = [$"torch{torchVersion}", $"torchvision{torchvisionVersion}"], | ||
| SkipTorchInstall = true, | ||
| }; | ||
|
|
||
| await StandardPipInstallProcessAsync( | ||
| venvRunner, | ||
| options, | ||
| installedPackage, | ||
| config, | ||
| onConsoleOutput, | ||
| progress, | ||
| cancellationToken | ||
| ) | ||
| .ConfigureAwait(false); | ||
|
|
||
| venvRunner.UpdateEnvironmentVariables(env => env.Remove("DISABLE_CUDA")); | ||
|
Member
There was a problem hiding this comment. nit: this |
||
| } | ||
|
|
||
| public override async Task RunPackage( | ||
| string installLocation, | ||
| InstalledPackage installedPackage, | ||
| RunPackageOptions options, | ||
| Action<ProcessOutput>? onConsoleOutput = null, | ||
| CancellationToken cancellationToken = default | ||
| ) | ||
| { | ||
| await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion)) | ||
| .ConfigureAwait(false); | ||
|
|
||
| // Desktop Tkinter app - there is no local URL to wait for, so startup is complete | ||
| // as soon as the process is up. | ||
| VenvRunner.RunDetached( | ||
| [Path.Combine(installLocation, options.Command ?? LaunchCommand), .. options.Arguments], | ||
| onConsoleOutput, | ||
| OnExit | ||
| ); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Totally optional and fine as is: this switch just repeats each package's
DisplayName, so a_ => args.CurrentPackagePair.BasePackage.DisplayNamedefault would remove the need for this arm (and fix packages like AI-Toolkit showing "Running Stable Diffusion"). Happy to leave that for a follow-up on our side.