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 virtual bool IsInferenceCompatible => false;
public abstract string OutputFolderName { get; }
public abstract Task DownloadPackage(
string installLocation,
DownloadPackageVersionOptions versionOptions,
IProgress? progress1
);
public abstract Task InstallPackage(
string installLocation,
TorchVersion torchVersion,
DownloadPackageVersionOptions versionOptions,
IProgress? progress = null,
Action? onConsoleOutput = null
);
public abstract Task RunPackage(
string installedPackagePath,
string command,
string arguments,
Action? onConsoleOutput
);
public abstract Task CheckForUpdates(InstalledPackage package);
public abstract Task Update(
InstalledPackage installedPackage,
TorchVersion torchVersion,
DownloadPackageVersionOptions versionOptions,
IProgress? progress = null,
bool includePrerelease = false,
Action? onConsoleOutput = null
);
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 Task SetupOutputFolderLinks(DirectoryPath installDirectory);
public abstract Task RemoveOutputFolderLinks(DirectoryPath installDirectory);
public abstract IEnumerable AvailableTorchVersions { get; }
public virtual bool IsCompatible => GetRecommendedTorchVersion() != TorchVersion.Cpu;
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 abstract Dictionary>? SharedFolders { get; }
public abstract Dictionary<
SharedOutputType,
IReadOnlyList
>? SharedOutputFolders { 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? Exited;
public event EventHandler? StartupComplete;
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,
Action? onConsoleOutput = null
)
{
progress?.Report(
new ProgressReport(-1f, "Installing PyTorch for CUDA", isIndeterminate: true)
);
await venvRunner
.PipInstall(
new PipInstallArgs()
.WithTorch("==2.0.1")
.WithXFormers("==0.0.20")
.WithTorchExtraIndex("cu118"),
onConsoleOutput
)
.ConfigureAwait(false);
}
protected Task InstallDirectMlTorch(
PyVenvRunner venvRunner,
IProgress? progress = null,
Action? onConsoleOutput = null
)
{
progress?.Report(
new ProgressReport(-1f, "Installing PyTorch for DirectML", isIndeterminate: true)
);
return venvRunner.PipInstall(new PipInstallArgs().WithTorchDirectML(), onConsoleOutput);
}
protected Task InstallCpuTorch(
PyVenvRunner venvRunner,
IProgress? progress = null,
Action? onConsoleOutput = null
)
{
progress?.Report(
new ProgressReport(-1f, "Installing PyTorch for CPU", isIndeterminate: true)
);
return venvRunner.PipInstall(
new PipInstallArgs().WithTorch("==2.0.1").WithTorchVision(),
onConsoleOutput
);
}
}