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.

338 lines
13 KiB

using System.Diagnostics;
using System.Drawing;
1 year ago
using System.Text.RegularExpressions;
using NLog;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Helper.Cache;
using StabilityMatrix.Core.Models.FileInterfaces;
using StabilityMatrix.Core.Models.Progress;
using StabilityMatrix.Core.Processes;
using StabilityMatrix.Core.Python;
using StabilityMatrix.Core.Services;
using YamlDotNet.Serialization;
using YamlDotNet.Serialization.NamingConventions;
1 year ago
namespace StabilityMatrix.Core.Models.Packages;
1 year ago
public class ComfyUI : BaseGitPackage
{
private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
public override string Name => "ComfyUI";
public override string DisplayName { get; set; } = "ComfyUI";
public override string Author => "comfyanonymous";
public override string LicenseType => "GPL-3.0";
1 year ago
public override string LicenseUrl =>
"https://github.com/comfyanonymous/ComfyUI/blob/master/LICENSE";
public override string Blurb => "A powerful and modular stable diffusion GUI and backend";
1 year ago
public override string LaunchCommand => "main.py";
public override Uri PreviewImageUri =>
new("https://github.com/comfyanonymous/ComfyUI/raw/master/comfyui_screenshot.png");
public override bool ShouldIgnoreReleases => true;
public override SharedFolderMethod RecommendedSharedFolderMethod =>
SharedFolderMethod.Configuration;
1 year ago
1 year ago
public ComfyUI(
IGithubApiCache githubApi,
ISettingsManager settingsManager,
IDownloadService downloadService,
IPrerequisiteHelper prerequisiteHelper
)
: base(githubApi, settingsManager, downloadService, prerequisiteHelper) { }
1 year ago
// https://github.com/comfyanonymous/ComfyUI/blob/master/folder_paths.py#L11
1 year ago
public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders =>
1 year ago
new()
{
1 year ago
[SharedFolderType.StableDiffusion] = new[] { "models/checkpoints" },
[SharedFolderType.Diffusers] = new[] { "models/diffusers" },
[SharedFolderType.Lora] = new[] { "models/loras" },
[SharedFolderType.CLIP] = new[] { "models/clip" },
[SharedFolderType.TextualInversion] = new[] { "models/embeddings" },
[SharedFolderType.VAE] = new[] { "models/vae" },
[SharedFolderType.ApproxVAE] = new[] { "models/vae_approx" },
[SharedFolderType.ControlNet] = new[] { "models/controlnet" },
[SharedFolderType.GLIGEN] = new[] { "models/gligen" },
[SharedFolderType.ESRGAN] = new[] { "models/upscale_models" },
[SharedFolderType.Hypernetwork] = new[] { "models/hypernetworks" },
};
public override List<LaunchOptionDefinition> LaunchOptions =>
new List<LaunchOptionDefinition>
{
new()
{
1 year ago
Name = "VRAM",
Type = LaunchOptionType.Bool,
InitialValue = HardwareHelper
.IterGpuInfo()
.Select(gpu => gpu.MemoryLevel)
.Max() switch
{
Level.Low => "--lowvram",
Level.Medium => "--normalvram",
_ => null
},
Options = { "--highvram", "--normalvram", "--lowvram", "--novram" }
},
1 year ago
new()
{
Name = "Enable DirectML",
Type = LaunchOptionType.Bool,
InitialValue = HardwareHelper.PreferDirectML(),
Options = { "--directml" }
},
new()
1 year ago
{
Name = "Use CPU only",
Type = LaunchOptionType.Bool,
InitialValue = !HardwareHelper.HasNvidiaGpu() && !HardwareHelper.HasAmdGpu(),
1 year ago
Options = { "--cpu" }
},
new()
{
Name = "Disable Xformers",
Type = LaunchOptionType.Bool,
InitialValue = !HardwareHelper.HasNvidiaGpu(),
Options = { "--disable-xformers" }
},
new()
{
Name = "Auto-Launch",
Type = LaunchOptionType.Bool,
Options = { "--auto-launch" }
},
LaunchOptionDefinition.Extras
};
1 year ago
public override Task<string> GetLatestVersion() => Task.FromResult("master");
1 year ago
public override IEnumerable<TorchVersion> AvailableTorchVersions =>
new[] { TorchVersion.Cpu, TorchVersion.Cuda, TorchVersion.DirectMl, TorchVersion.Rocm };
public override async Task InstallPackage(
string installLocation,
TorchVersion torchVersion,
IProgress<ProgressReport>? progress = null,
Action<ProcessOutput>? onConsoleOutput = null
1 year ago
)
1 year ago
{
await base.InstallPackage(installLocation, torchVersion, progress).ConfigureAwait(false);
progress?.Report(new ProgressReport(-1, "Setting up venv", isIndeterminate: true));
// Setup venv
await using var venvRunner = new PyVenvRunner(Path.Combine(installLocation, "venv"));
venvRunner.WorkingDirectory = installLocation;
await venvRunner.Setup(true).ConfigureAwait(false);
// Install torch / xformers based on gpu info
switch (torchVersion)
{
case TorchVersion.Cpu:
await InstallCpuTorch(venvRunner, progress, onConsoleOutput).ConfigureAwait(false);
break;
case TorchVersion.Cuda:
await InstallCudaTorch(venvRunner, progress, onConsoleOutput).ConfigureAwait(false);
break;
case TorchVersion.Rocm:
await InstallRocmTorch(venvRunner, progress, onConsoleOutput).ConfigureAwait(false);
break;
case TorchVersion.DirectMl:
await InstallDirectMlTorch(venvRunner, progress, onConsoleOutput)
.ConfigureAwait(false);
break;
default:
throw new ArgumentOutOfRangeException(nameof(torchVersion), torchVersion, null);
}
// Install requirements file
1 year ago
progress?.Report(
new ProgressReport(-1, "Installing Package Requirements", isIndeterminate: true)
);
Logger.Info("Installing requirements.txt");
await venvRunner.PipInstall($"-r requirements.txt", onConsoleOutput).ConfigureAwait(false);
1 year ago
progress?.Report(
new ProgressReport(1, "Installing Package Requirements", isIndeterminate: false)
);
}
1 year ago
private async Task AutoDetectAndInstallTorch(
PyVenvRunner venvRunner,
IProgress<ProgressReport>? progress = null
)
{
var gpus = HardwareHelper.IterGpuInfo().ToList();
if (gpus.Any(g => g.IsNvidia))
{
await InstallCudaTorch(venvRunner, progress).ConfigureAwait(false);
}
else if (HardwareHelper.PreferRocm())
{
await InstallRocmTorch(venvRunner, progress).ConfigureAwait(false);
}
else if (HardwareHelper.PreferDirectML())
{
await InstallDirectMlTorch(venvRunner, progress).ConfigureAwait(false);
}
else
{
await InstallCpuTorch(venvRunner, progress).ConfigureAwait(false);
}
1 year ago
}
1 year ago
public override async Task RunPackage(
string installedPackagePath,
string command,
string arguments,
Action<ProcessOutput>? onConsoleOutput
1 year ago
)
1 year ago
{
await SetupVenv(installedPackagePath).ConfigureAwait(false);
1 year ago
void HandleConsoleOutput(ProcessOutput s)
1 year ago
{
onConsoleOutput?.Invoke(s);
1 year ago
if (s.Text.Contains("To see the GUI go to", StringComparison.OrdinalIgnoreCase))
1 year ago
{
var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
var match = regex.Match(s.Text);
1 year ago
if (match.Success)
{
WebUrl = match.Value;
}
OnStartupComplete(WebUrl);
}
}
void HandleExit(int i)
{
Debug.WriteLine($"Venv process exited with code {i}");
OnExit(i);
}
var args = $"\"{Path.Combine(installedPackagePath, command)}\" {arguments}";
1 year ago
1 year ago
VenvRunner?.RunDetached(args.TrimEnd(), HandleConsoleOutput, HandleExit);
1 year ago
}
1 year ago
public override Task SetupModelFolders(
DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod
)
{
switch (sharedFolderMethod)
{
case SharedFolderMethod.None:
return Task.CompletedTask;
case SharedFolderMethod.Symlink:
return base.SetupModelFolders(installDirectory, sharedFolderMethod);
}
var extraPathsYamlPath = installDirectory + "extra_model_paths.yaml";
var modelsDir = SettingsManager.ModelsDirectory;
var deserializer = new DeserializerBuilder()
.WithNamingConvention(UnderscoredNamingConvention.Instance)
.IgnoreUnmatchedProperties()
.Build();
var exists = File.Exists(extraPathsYamlPath);
if (!exists)
{
Logger.Info("Creating extra_model_paths.yaml");
File.WriteAllText(extraPathsYamlPath, string.Empty);
}
var yaml = File.ReadAllText(extraPathsYamlPath);
1 year ago
var comfyModelPaths =
deserializer.Deserialize<ComfyModelPathsYaml>(yaml)
??
// ReSharper disable once NullCoalescingConditionIsAlwaysNotNullAccordingToAPIContract
// cuz it can actually be null lol
new ComfyModelPathsYaml();
comfyModelPaths.StabilityMatrix ??= new ComfyModelPathsYaml.SmData();
comfyModelPaths.StabilityMatrix.Checkpoints = Path.Combine(modelsDir, "StableDiffusion");
comfyModelPaths.StabilityMatrix.Vae = Path.Combine(modelsDir, "VAE");
1 year ago
comfyModelPaths.StabilityMatrix.Loras =
$"{Path.Combine(modelsDir, "Lora")}\n" + $"{Path.Combine(modelsDir, "LyCORIS")}";
comfyModelPaths.StabilityMatrix.UpscaleModels =
$"{Path.Combine(modelsDir, "ESRGAN")}\n"
+ $"{Path.Combine(modelsDir, "RealESRGAN")}\n"
+ $"{Path.Combine(modelsDir, "SwinIR")}";
comfyModelPaths.StabilityMatrix.Embeddings = Path.Combine(modelsDir, "TextualInversion");
comfyModelPaths.StabilityMatrix.Hypernetworks = Path.Combine(modelsDir, "Hypernetwork");
comfyModelPaths.StabilityMatrix.Controlnet = Path.Combine(modelsDir, "ControlNet");
comfyModelPaths.StabilityMatrix.Clip = Path.Combine(modelsDir, "CLIP");
comfyModelPaths.StabilityMatrix.Diffusers = Path.Combine(modelsDir, "Diffusers");
comfyModelPaths.StabilityMatrix.Gligen = Path.Combine(modelsDir, "GLIGEN");
comfyModelPaths.StabilityMatrix.VaeApprox = Path.Combine(modelsDir, "ApproxVAE");
1 year ago
var serializer = new SerializerBuilder()
.WithNamingConvention(UnderscoredNamingConvention.Instance)
.Build();
var yamlData = serializer.Serialize(comfyModelPaths);
File.WriteAllText(extraPathsYamlPath, yamlData);
return Task.CompletedTask;
}
1 year ago
public override Task UpdateModelFolders(
DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod
) => SetupModelFolders(installDirectory, sharedFolderMethod);
1 year ago
public override Task RemoveModelFolderLinks(
DirectoryPath installDirectory,
SharedFolderMethod sharedFolderMethod
)
{
return sharedFolderMethod switch
{
SharedFolderMethod.Configuration => Task.CompletedTask,
SharedFolderMethod.None => Task.CompletedTask,
1 year ago
SharedFolderMethod.Symlink
=> base.RemoveModelFolderLinks(installDirectory, sharedFolderMethod),
_ => Task.CompletedTask
};
}
1 year ago
private async Task InstallRocmTorch(
PyVenvRunner venvRunner,
IProgress<ProgressReport>? progress = null,
Action<ProcessOutput>? onConsoleOutput = null
1 year ago
)
{
1 year ago
progress?.Report(
new ProgressReport(-1f, "Installing PyTorch for ROCm", isIndeterminate: true)
);
await venvRunner.PipInstall("--upgrade pip wheel", onConsoleOutput).ConfigureAwait(false);
1 year ago
await venvRunner
.PipInstall(PyVenvRunner.TorchPipInstallArgsRocm542, onConsoleOutput)
.ConfigureAwait(false);
}
public class ComfyModelPathsYaml
{
public class SmData
{
public string Checkpoints { get; set; }
public string Vae { get; set; }
public string Loras { get; set; }
public string UpscaleModels { get; set; }
public string Embeddings { get; set; }
public string Hypernetworks { get; set; }
public string Controlnet { get; set; }
public string Clip { get; set; }
public string Diffusers { get; set; }
public string Gligen { get; set; }
public string VaeApprox { get; set; }
}
public SmData? StabilityMatrix { get; set; }
}
1 year ago
}