Merge pull request #565 from ionite34/diy-backport

Diy backport
LykosAI · Mar 13, 2024 · f2051df · f2051df
2 parents a09775a + e4aa0f8
commit f2051df
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 16 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - Fixed [#502](https://github.com/LykosAI/StabilityMatrix/issues/502) - missing launch options for Forge
 - Fixed [#500](https://github.com/LykosAI/StabilityMatrix/issues/500) - missing output images in Forge when using output sharing
 - Fixed [#490](https://github.com/LykosAI/StabilityMatrix/issues/490) - `mpmath has no attribute 'rational'` error on macOS
+- Fixed [#510](https://github.com/ionite34/StabilityMatrix/pull/564/files) - kohya_ss packages with v23.0.x failing to install due to missing 'packaging' dependency
 - Fixed incorrect progress text when deleting a checkpoint from the Checkpoints page
 - Fixed incorrect icon colors on macOS
 

diff --git a/StabilityMatrix.Core/Models/Packages/KohyaSs.cs b/StabilityMatrix.Core/Models/Packages/KohyaSs.cs
@@ -116,12 +116,24 @@ public override async Task InstallPackage(
         Action<ProcessOutput>? onConsoleOutput = null
     )
     {
+        progress?.Report(new ProgressReport(-1f, "Updating submodules", isIndeterminate: true));
+        await PrerequisiteHelper
+            .RunGit(
+                ["submodule", "update", "--init", "--recursive", "--quiet"],
+                onConsoleOutput,
+                installLocation
+            )
+            .ConfigureAwait(false);
+
         progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true));
         // Setup venv
         await using var venvRunner = new PyVenvRunner(Path.Combine(installLocation, "venv"));
         venvRunner.WorkingDirectory = installLocation;
         await venvRunner.Setup(true, onConsoleOutput).ConfigureAwait(false);
 
+        // Extra dep needed before running setup since v23.0.x
+        await venvRunner.PipInstall("packaging").ConfigureAwait(false);
+
         if (Compat.IsWindows)
         {
             var setupSmPath = Path.Combine(installLocation, "setup", "setup_sm.py");
@@ -136,18 +148,17 @@ import setup_common
             await File.WriteAllTextAsync(setupSmPath, setupText).ConfigureAwait(false);
 
             // Install
-            venvRunner.RunDetached("setup/setup_sm.py", onConsoleOutput);
-            await venvRunner.Process.WaitForExitAsync().ConfigureAwait(false);
-
+            await venvRunner.CustomInstall("setup/setup_sm.py", onConsoleOutput).ConfigureAwait(false);
             await venvRunner.PipInstall("bitsandbytes-windows").ConfigureAwait(false);
         }
         else if (Compat.IsLinux)
         {
-            venvRunner.RunDetached(
-                "setup/setup_linux.py --platform-requirements-file=requirements_linux.txt --no_run_accelerate",
-                onConsoleOutput
-            );
-            await venvRunner.Process.WaitForExitAsync().ConfigureAwait(false);
+            await venvRunner
+                .CustomInstall(
+                    "setup/setup_linux.py --platform-requirements-file=requirements_linux.txt --no_run_accelerate",
+                    onConsoleOutput
+                )
+                .ConfigureAwait(false);
         }
     }
 
@@ -198,18 +209,40 @@ def rewrite_module(self, module_text: str) -> str:
                 1.ToPython()
             );
 
-            var filesToUpdate = new[]
+            var kohyaGuiDir = Path.Combine(installedPackagePath, "kohya_gui");
+            var guiDirExists = Directory.Exists(kohyaGuiDir);
+            var filesToUpdate = new List<string>();
+            if (guiDirExists)
             {
-                "lora_gui.py",
-                "dreambooth_gui.py",
-                "textual_inversion_gui.py",
-                Path.Combine("library", "wd14_caption_gui.py"),
-                "finetune_gui.py"
-            };
+                filesToUpdate.AddRange(
+                    [
+                        "lora_gui.py",
+                        "dreambooth_gui.py",
+                        "textual_inversion_gui.py",
+                        "wd14_caption_gui.py",
+                        "finetune_gui.py"
+                    ]
+                );
+            }
+            else
+            {
+                filesToUpdate.AddRange(
+                    [
+                        "lora_gui.py",
+                        "dreambooth_gui.py",
+                        "textual_inversion_gui.py",
+                        Path.Combine("library", "wd14_caption_gui.py"),
+                        "finetune_gui.py"
+                    ]
+                );
+            }
 
             foreach (var file in filesToUpdate)
             {
-                var path = Path.Combine(installedPackagePath, file);
+                var path = Path.Combine(guiDirExists ? kohyaGuiDir : installedPackagePath, file);
+                if (!File.Exists(path))
+                    continue;
+
                 var text = File.ReadAllText(path);
                 if (text.Contains(replacementAcceleratePath.Replace(@"\", @"\\")))
                     continue;