Multi-Platform Package Manager for Stable Diffusion
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

409 lines
15 KiB

using System.Globalization;
using System.Text.RegularExpressions;
using NLog;
using StabilityMatrix.Core.Attributes;
using StabilityMatrix.Core.Extensions;
1 year ago
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Helper.Cache;
using StabilityMatrix.Core.Models.FileInterfaces;
1 year ago
using StabilityMatrix.Core.Models.Progress;
using StabilityMatrix.Core.Processes;
1 year ago
using StabilityMatrix.Core.Python;
using StabilityMatrix.Core.Services;
namespace StabilityMatrix.Core.Models.Packages;
[Singleton(typeof(BasePackage))]
1 year ago
public class InvokeAI : BaseGitPackage
{
private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
private const string RelativeRootPath = "invokeai-root";
1 year ago
public override string Name => "InvokeAI";
public override string DisplayName { get; set; } = "InvokeAI";
public override string Author => "invoke-ai";
public override string LicenseType => "Apache-2.0";
1 year ago
public override string LicenseUrl => "https://github.com/invoke-ai/InvokeAI/blob/main/LICENSE";
1 year ago
public override string Blurb => "Professional Creative Tools for Stable Diffusion";
public override string LaunchCommand => "invokeai-web";
public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Nightmare;
1 year ago
public override IReadOnlyList<string> ExtraLaunchCommands =>
new[]
{
"invokeai-configure",
"invokeai-merge",
"invokeai-metadata",
"invokeai-model-install",
"invokeai-node-cli",
"invokeai-ti",
"invokeai-update",
};
public override Uri PreviewImageUri =>
new(
"https://raw.githubusercontent.com/invoke-ai/InvokeAI/main/docs/assets/canvas_preview.png"
);
1 year ago
public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods =>
new[] { SharedFolderMethod.Symlink, SharedFolderMethod.None };
public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.Symlink;
public override string MainBranch => "main";
1 year ago
public InvokeAI(
IGithubApiCache githubApi,
ISettingsManager settingsManager,
IDownloadService downloadService,
1 year ago
IPrerequisiteHelper prerequisiteHelper
)
: base(githubApi, settingsManager, downloadService, prerequisiteHelper) { }
1 year ago
public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders =>
new()
{
[SharedFolderType.StableDiffusion] = new[]
{
Path.Combine(RelativeRootPath, "autoimport", "main")
},
[SharedFolderType.Lora] = new[]
{
Path.Combine(RelativeRootPath, "autoimport", "lora")
},
1 year ago
[SharedFolderType.TextualInversion] = new[]
{
Path.Combine(RelativeRootPath, "autoimport", "embedding")
1 year ago
},
[SharedFolderType.ControlNet] = new[]
{
Path.Combine(RelativeRootPath, "autoimport", "controlnet")
},
[SharedFolderType.IpAdapter] = new[]
{
Path.Combine(RelativeRootPath, "autoimport", "ip_adapter")
}
1 year ago
};
public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders =>
new()
{
[SharedOutputType.Text2Img] = new[]
{
Path.Combine("invokeai-root", "outputs", "images")
}
};
public override string OutputFolderName => Path.Combine("invokeai-root", "outputs", "images");
1 year ago
// https://github.com/invoke-ai/InvokeAI/blob/main/docs/features/CONFIGURATION.md
public override List<LaunchOptionDefinition> LaunchOptions =>
new List<LaunchOptionDefinition>
{
1 year ago
new()
{
Name = "Host",
Type = LaunchOptionType.String,
DefaultValue = "localhost",
Options = new List<string> { "--host" }
},
new()
{
Name = "Port",
Type = LaunchOptionType.String,
DefaultValue = "9090",
Options = new List<string> { "--port" }
},
new()
{
Name = "Allow Origins",
Description =
"List of host names or IP addresses that are allowed to connect to the "
+ "InvokeAI API in the format ['host1','host2',...]",
Type = LaunchOptionType.String,
DefaultValue = "[]",
Options = new List<string> { "--allow-origins" }
},
new()
{
Name = "Always use CPU",
Type = LaunchOptionType.Bool,
Options = new List<string> { "--always_use_cpu" }
},
new()
{
Name = "Precision",
Type = LaunchOptionType.Bool,
Options = new List<string>
{
"--precision auto",
"--precision float16",
"--precision float32",
}
},
new()
{
Name = "Aggressively free up GPU memory after each operation",
Type = LaunchOptionType.Bool,
Options = new List<string> { "--free_gpu_mem" }
},
LaunchOptionDefinition.Extras
};
1 year ago
public override IEnumerable<TorchVersion> AvailableTorchVersions =>
new[] { TorchVersion.Cpu, TorchVersion.Cuda, TorchVersion.Rocm, TorchVersion.Mps };
public override TorchVersion GetRecommendedTorchVersion()
{
if (Compat.IsMacOS && Compat.IsArm)
{
return TorchVersion.Mps;
}
return base.GetRecommendedTorchVersion();
}
1 year ago
public override async Task InstallPackage(
string installLocation,
TorchVersion torchVersion,
SharedFolderMethod selectedSharedFolderMethod,
DownloadPackageVersionOptions versionOptions,
IProgress<ProgressReport>? progress = null,
Action<ProcessOutput>? onConsoleOutput = null
1 year ago
)
1 year ago
{
// Setup venv
progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true));
var venvPath = Path.Combine(installLocation, "venv");
var exists = Directory.Exists(venvPath);
await using var venvRunner = new PyVenvRunner(venvPath);
venvRunner.WorkingDirectory = installLocation;
await venvRunner.Setup(true, onConsoleOutput).ConfigureAwait(false);
venvRunner.EnvironmentVariables = GetEnvVars(installLocation);
progress?.Report(new ProgressReport(-1f, "Installing Package", isIndeterminate: true));
var pipCommandArgs =
"-e . --use-pep517 --extra-index-url https://download.pytorch.org/whl/cpu";
switch (torchVersion)
1 year ago
{
// If has Nvidia Gpu, install CUDA version
case TorchVersion.Cuda:
progress?.Report(
new ProgressReport(-1f, "Installing PyTorch for CUDA", isIndeterminate: true)
);
var args = new List<Argument>();
if (exists)
{
var pipPackages = await venvRunner.PipList().ConfigureAwait(false);
var hasCuda121 = pipPackages.Any(
p => p.Name == "torch" && p.Version.Contains("cu121")
);
if (!hasCuda121)
{
args.Add("--upgrade");
args.Add("--force-reinstall");
}
}
await venvRunner
.PipInstall(
new PipInstallArgs(args.Any() ? args.ToArray() : Array.Empty<Argument>())
.WithTorch("==2.1.0")
.WithTorchVision("==0.16.0")
.WithXFormers("==0.0.22post7")
.WithTorchExtraIndex("cu121"),
onConsoleOutput
)
.ConfigureAwait(false);
Logger.Info("Starting InvokeAI install (CUDA)...");
pipCommandArgs =
"-e .[xformers] --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu121";
break;
// For AMD, Install ROCm version
case TorchVersion.Rocm:
await venvRunner
.PipInstall(
new PipInstallArgs()
.WithTorch("==2.0.1")
.WithTorchVision()
.WithExtraIndex("rocm5.4.2"),
onConsoleOutput
)
.ConfigureAwait(false);
Logger.Info("Starting InvokeAI install (ROCm)...");
pipCommandArgs =
"-e . --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.4.2";
break;
case TorchVersion.Mps:
// For Apple silicon, use MPS
Logger.Info("Starting InvokeAI install (MPS)...");
pipCommandArgs = "-e . --use-pep517";
break;
1 year ago
}
await venvRunner
.PipInstall($"{pipCommandArgs}{(exists ? " --upgrade" : "")}", onConsoleOutput)
.ConfigureAwait(false);
1 year ago
await venvRunner
.PipInstall("rich packaging python-dotenv", onConsoleOutput)
1 year ago
.ConfigureAwait(false);
progress?.Report(new ProgressReport(-1f, "Configuring InvokeAI", isIndeterminate: true));
// need to setup model links before running invokeai-configure so it can do its conversion
await SetupModelFolders(installLocation, selectedSharedFolderMethod).ConfigureAwait(false);
1 year ago
await RunInvokeCommand(
installLocation,
"invokeai-configure",
"--yes --skip-sd-weights",
true,
onConsoleOutput,
spam3: true
1 year ago
)
.ConfigureAwait(false);
await VenvRunner.Process.WaitForExitAsync();
progress?.Report(new ProgressReport(1f, "Done!", isIndeterminate: false));
1 year ago
}
1 year ago
public override Task RunPackage(
string installedPackagePath,
string command,
string arguments,
Action<ProcessOutput>? onConsoleOutput
) => RunInvokeCommand(installedPackagePath, command, arguments, true, onConsoleOutput);
1 year ago
private async Task RunInvokeCommand(
string installedPackagePath,
string command,
string arguments,
bool runDetached,
Action<ProcessOutput>? onConsoleOutput,
bool spam3 = false
1 year ago
)
1 year ago
{
if (spam3 && !runDetached)
{
throw new InvalidOperationException("Cannot spam 3 if not running detached");
}
1 year ago
await SetupVenv(installedPackagePath).ConfigureAwait(false);
arguments = command switch
{
"invokeai-configure" => "--yes --skip-sd-weights",
_ => arguments
};
VenvRunner.EnvironmentVariables = GetEnvVars(installedPackagePath);
// Launch command is for a console entry point, and not a direct script
var entryPoint = await VenvRunner.GetEntryPoint(command).ConfigureAwait(false);
// Split at ':' to get package and function
var split = entryPoint?.Split(':');
1 year ago
if (split is not { Length: > 1 })
{
throw new Exception($"Could not find entry point for InvokeAI: {entryPoint.ToRepr()}");
}
1 year ago
// Compile a startup command according to
// https://packaging.python.org/en/latest/specifications/entry-points/#use-for-scripts
// For invokeai, also patch the shutil.get_terminal_size function to return a fixed value
// above the minimum in invokeai.frontend.install.widgets
var code = $"""
try:
import os
import shutil
from invokeai.frontend.install import widgets
_min_cols = widgets.MIN_COLS
_min_lines = widgets.MIN_LINES
static_size_fn = lambda: os.terminal_size((_min_cols, _min_lines))
shutil.get_terminal_size = static_size_fn
widgets.get_terminal_size = static_size_fn
except Exception as e:
import warnings
warnings.warn('Could not patch terminal size for InvokeAI' + str(e))
import sys
from {split[0]} import {split[1]}
sys.exit({split[1]}())
""";
if (runDetached)
{
void HandleConsoleOutput(ProcessOutput s)
{
onConsoleOutput?.Invoke(s);
if (
spam3
&& s.Text.Contains(
"[3] Accept the best guess;",
StringComparison.OrdinalIgnoreCase
)
)
{
VenvRunner.Process?.StandardInput.WriteLine("3");
return;
}
1 year ago
if (!s.Text.Contains("running on", StringComparison.OrdinalIgnoreCase))
return;
1 year ago
var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
var match = regex.Match(s.Text);
if (!match.Success)
return;
1 year ago
WebUrl = match.Value;
OnStartupComplete(WebUrl);
}
1 year ago
VenvRunner.RunDetached(
$"-c \"{code}\" {arguments}".TrimEnd(),
HandleConsoleOutput,
OnExit
);
}
else
{
1 year ago
var result = await VenvRunner
.Run($"-c \"{code}\" {arguments}".TrimEnd())
.ConfigureAwait(false);
onConsoleOutput?.Invoke(new ProcessOutput { Text = result.StandardOutput });
}
1 year ago
}
private Dictionary<string, string> GetEnvVars(DirectoryPath installPath)
{
// Set additional required environment variables
var env = new Dictionary<string, string>();
if (SettingsManager.Settings.EnvironmentVariables is not null)
{
env.Update(SettingsManager.Settings.EnvironmentVariables);
}
// Need to make subdirectory because they store config in the
// directory *above* the root directory
var root = installPath.JoinDir(RelativeRootPath);
root.Create();
env["INVOKEAI_ROOT"] = root;
return env;
}
1 year ago
}