using Octokit;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Models.Database;
using StabilityMatrix.Core.Models.FileInterfaces;
using StabilityMatrix.Core.Models.Progress;
using StabilityMatrix.Core.Processes;
using StabilityMatrix.Core.Python;
namespace StabilityMatrix.Core.Models.Packages;
public abstract class BasePackage
{
public string ByAuthor => $"By {Author}";
public abstract string Name { get; }
public abstract string DisplayName { get; set; }
public abstract string Author { get; }
public abstract string Blurb { get; }
public abstract string GithubUrl { get; }
public abstract string LicenseType { get; }
public abstract string LicenseUrl { get; }
public virtual string Disclaimer => string.Empty;
public virtual bool OfferInOneClickInstaller => true;
///
/// Primary command to launch the package. 'Launch' buttons uses this.
///
public abstract string LaunchCommand { get; }
///
/// Optional commands (e.g. 'config') that are on the launch button split drop-down.
///
public virtual IReadOnlyList ExtraLaunchCommands { get; } = Array.Empty();
public abstract Uri PreviewImageUri { get; }
public virtual bool ShouldIgnoreReleases => false;
public virtual bool UpdateAvailable { get; set; }
public abstract Task DownloadPackage(string installLocation, DownloadPackageVersionOptions versionOptions,
IProgress? progress1);
public abstract Task InstallPackage(string installLocation, TorchVersion torchVersion,
IProgress? progress = null);
public abstract Task RunPackage(string installedPackagePath, string command, string arguments);
public abstract Task CheckForUpdates(InstalledPackage package);
public abstract Task Update(InstalledPackage installedPackage,
TorchVersion torchVersion, IProgress? progress = null,
bool includePrerelease = false);
public virtual IEnumerable AvailableSharedFolderMethods => new[]
{
SharedFolderMethod.Symlink,
SharedFolderMethod.Configuration,
SharedFolderMethod.None
};
public abstract SharedFolderMethod RecommendedSharedFolderMethod { get; }
public abstract Task SetupModelFolders(DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod);
public abstract Task UpdateModelFolders(DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod);
public abstract Task RemoveModelFolderLinks(DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod);
public abstract IEnumerable AvailableTorchVersions { get; }
public virtual TorchVersion GetRecommendedTorchVersion()
{
// if there's only one AvailableTorchVersion, return that
if (AvailableTorchVersions.Count() == 1)
{
return AvailableTorchVersions.First();
}
if (HardwareHelper.HasNvidiaGpu() && AvailableTorchVersions.Contains(TorchVersion.Cuda))
{
return TorchVersion.Cuda;
}
if (HardwareHelper.PreferRocm() && AvailableTorchVersions.Contains(TorchVersion.Rocm))
{
return TorchVersion.Rocm;
}
if (HardwareHelper.PreferDirectML() &&
AvailableTorchVersions.Contains(TorchVersion.DirectMl))
{
return TorchVersion.DirectMl;
}
return TorchVersion.Cpu;
}
///
/// Shuts down the subprocess, canceling any pending streams.
///
public abstract void Shutdown();
///
/// Shuts down the process, returning a Task to wait for output EOF.
///
public abstract Task WaitForShutdown();
public abstract Task> GetReleaseTags();
public abstract List LaunchOptions { get; }
public virtual string? ExtraLaunchArguments { get; set; } = null;
///
/// The shared folders that this package supports.
/// Mapping of to the relative paths from the package root.
///
public virtual Dictionary>? SharedFolders { get; }
public abstract Task GetLatestVersion();
public abstract Task GetAllVersionOptions();
public abstract Task?> GetAllCommits(string branch, int page = 1, int perPage = 10);
public abstract Task> GetAllBranches();
public abstract Task> GetAllReleases();
public event EventHandler? ConsoleOutput;
public event EventHandler? Exited;
public event EventHandler? StartupComplete;
public void OnConsoleOutput(ProcessOutput output) => ConsoleOutput?.Invoke(this, output);
public void OnExit(int exitCode) => Exited?.Invoke(this, exitCode);
public void OnStartupComplete(string url) => StartupComplete?.Invoke(this, url);
public virtual PackageVersionType AvailableVersionTypes => ShouldIgnoreReleases
? PackageVersionType.Commit
: PackageVersionType.GithubRelease | PackageVersionType.Commit;
protected async Task InstallCudaTorch(PyVenvRunner venvRunner,
IProgress? progress = null)
{
progress?.Report(new ProgressReport(-1f, "Installing PyTorch for CUDA",
isIndeterminate: true));
await venvRunner.PipInstall(PyVenvRunner.TorchPipInstallArgsCuda, OnConsoleOutput)
.ConfigureAwait(false);
await venvRunner.PipInstall("xformers", OnConsoleOutput).ConfigureAwait(false);
}
protected async Task InstallDirectMlTorch(PyVenvRunner venvRunner,
IProgress? progress = null)
{
progress?.Report(new ProgressReport(-1f, "Installing PyTorch for DirectML",
isIndeterminate: true));
await venvRunner.PipInstall(PyVenvRunner.TorchPipInstallArgsDirectML, OnConsoleOutput)
.ConfigureAwait(false);
}
protected async Task InstallCpuTorch(PyVenvRunner venvRunner, IProgress? progress = null)
{
progress?.Report(new ProgressReport(-1f, "Installing PyTorch for CPU",
isIndeterminate: true));
await venvRunner.PipInstall(PyVenvRunner.TorchPipInstallArgsCpu, OnConsoleOutput)
.ConfigureAwait(false);
}
}