diff --git a/CHANGELOG.md b/CHANGELOG.md index b15507a6..4363ff89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,17 +5,99 @@ All notable changes to Stability Matrix will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2.0.0.html). -## v2.10.0-dev.1 + +## v2.10.0-preview.1 ### Added - Added OpenArt.AI workflow browser for ComfyUI workflows +## v2.10.0-dev.3 +### Added +- Added support for deep links from the new Stability Matrix Chrome extension +### Changed +- Due to changes on the CivitAI API, you can no longer select a specific page in the CivitAI Model Browser +- Due to the above API changes, new pages are now loaded via "infinite scrolling" +### Fixed +- Fixed Inference HiresFix module "Inherit Primary Sampler Addons" setting not effectively disabling when unchecked +- Fixed model download location options for VAEs in the CivitAI Model Browser +- Fixed crash on startup when library directory is not set +- Fixed One-Click install progress dialog not disappearing after completion +- Fixed ComfyUI with Inference pop-up during one-click install appearing below the visible scroll area +- Fixed no packages being available for one-click install on PCs without a GPU +- Fixed models not being removed from the installed models cache when deleting them from the Checkpoints page +- Fixed missing ratings on some models in the CivitAI Model Browser +- Fixed missing favorite count in the CivitAI Model Browser +- Fixed recommended models not showing all SDXL models + +## v2.10.0-dev.2 +### Added +- Added Reference-Only mode for Inference ControlNet, used for guiding the sampler with an image without a pretrained model. Part of the latent and attention layers will be connected to the reference image, similar to Image to Image or Inpainting. +### Changed +- Inference Primary Sampler Addons (i.e. ControlNet, FreeU) are now inherited by Hires Fix Samplers, this can be overriden from the Hires Fix module's settings menu by disabling the "Inherit Primary Sampler Addons" option. +- Revisited the way images are loaded on the outputs page, with improvements to loading speed & not freezing the UI while loading +### Fixed +- Fixed Outputs page not remembering where the user last was in the TreeView in certain circumstances +- Fixed Inference extension upgrades not being added to missing extensions list for prompted install +- Fixed "The Open Web UI button has moved" teaching tip spam + +## v2.10.0-dev.1 +### Added +- Inference ControlNet module now supports over 42 preprocessors, a new button next to the preprocessors dropdown allows previewing the output of the selected preprocessor on the image. +- Added resolution selection for Inference ControlNet module, this controls preprocessor resolution too. +### Changed +- Revamped the Packages page to enable running multiple packages at the same time +- Changed the Outputs Page to use a TreeView for the directory selection instead of a dropdown selector +### Removed +- Removed the main Launch page, as it is no longer needed with the new Packages page + +## v2.9.1 +### Added +- Fixed [#498](https://github.com/LykosAI/StabilityMatrix/issues/498) Added "Pony" category to CivitAI Model Browser +### Changed +- Changed package deletion warning dialog to require additional confirmation +### Fixed +- 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 + +## v2.9.0 +### Added +- Added new package: [StableSwarmUI](https://github.com/Stability-AI/StableSwarmUI) by Stability AI +- Added new package: [Stable Diffusion WebUI Forge](https://github.com/lllyasviel/stable-diffusion-webui-forge) by lllyasviel +- Added extension management for SD.Next and Stable Diffusion WebUI-UX +- Added the ability to choose where CivitAI model downloads are saved +- Added `--launch-package` argument to launch a specific package on startup, using display name or package ID (i.e. `--launch-package "Stable Diffusion WebUI Forge"` or `--launch-package c0b3ecc5-9664-4be9-952d-a10b3dcaee14`) +- Added more Base Model search options to the CivitAI Model Browser +- Added Stable Cascade to the HuggingFace Model Browser +#### Inference +- Added Inference Prompt Styles, with Prompt Expansion model support (i.e. Fooocus V2) +- Added option to load a .yaml config file next to the model with the same name. Can be used with VPred and other models that require a config file. +- Added copy image support on linux and macOS for Inference outputs viewer menu +### Changed +- Updated translations for German, Spanish, French, Japanese, Portuguese, and Turkish +- (Internal) Updated to Avalonia 11.0.9 +### Fixed +- Fixed StableSwarmUI not installing properly on macOS +- Fixed [#464](https://github.com/LykosAI/StabilityMatrix/issues/464) - error when installing InvokeAI on macOS +- Fixed [#335](https://github.com/LykosAI/StabilityMatrix/issues/335) Update hanging indefinitely after git step for Auto1111 and SDWebUI Forge +- Fixed Inference output viewer menu "Copy" not copying image +- Fixed image viewer dialog arrow key navigation not working +- Fixed CivitAI login prompt not showing when downloading models that require CivitAI logins +- Fixed unknown model types not showing on checkpoints page (thanks Jerry!) +- Improved error handling for Inference Select Image hash calculation in case file is being written to while being read + ## v2.9.0-pre.2 ### Added - Added `--launch-package` argument to launch a specific package on startup, using display name or package ID (i.e. `--launch-package "Stable Diffusion WebUI Forge"` or `--launch-package c0b3ecc5-9664-4be9-952d-a10b3dcaee14`) +- Added more Base Model search options to the CivitAI Model Browser +- Added Stable Cascade to the HuggingFace Model Browser ### Changed - (Internal) Updated to Avalonia 11.0.9 ### Fixed - Fixed image viewer dialog arrow key navigation not working +- Fixed CivitAI login prompt not showing when downloading models that require CivitAI logins ## v2.9.0-pre.1 ### Added diff --git a/StabilityMatrix.Avalonia/App.axaml b/StabilityMatrix.Avalonia/App.axaml index 1f039a8a..fa669962 100644 --- a/StabilityMatrix.Avalonia/App.axaml +++ b/StabilityMatrix.Avalonia/App.axaml @@ -4,6 +4,7 @@ xmlns:local="using:StabilityMatrix.Avalonia" xmlns:idcr="using:Dock.Avalonia.Controls.Recycling" xmlns:styling="clr-namespace:FluentAvalonia.Styling;assembly=FluentAvalonia" + xmlns:controls="clr-namespace:Avalonia.Labs.Controls;assembly=Avalonia.Labs.Controls" Name="Stability Matrix" RequestedThemeVariant="Dark"> @@ -54,6 +55,7 @@ + @@ -79,6 +81,7 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/StabilityMatrix.Avalonia/Styles/ThemeColors.axaml b/StabilityMatrix.Avalonia/Styles/ThemeColors.axaml index dbca9287..963fb0f7 100644 --- a/StabilityMatrix.Avalonia/Styles/ThemeColors.axaml +++ b/StabilityMatrix.Avalonia/Styles/ThemeColors.axaml @@ -21,6 +21,7 @@ #00BCD4 #009688 #2C582C + #2C582C #3A783C #4BA04F #AA4BA04F diff --git a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs index 2d70b02e..d445488d 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs @@ -16,10 +16,9 @@ using Avalonia.Threading; using CommunityToolkit.Mvvm.Input; using ExifLibrary; using FluentAvalonia.UI.Controls; -using Microsoft.Extensions.DependencyInjection; -using Nito.Disposables.Internals; using NLog; using Refit; +using Semver; using SkiaSharp; using StabilityMatrix.Avalonia.Extensions; using StabilityMatrix.Avalonia.Helpers; @@ -59,6 +58,7 @@ public abstract partial class InferenceGenerationViewModelBase private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); private readonly ISettingsManager settingsManager; + private readonly RunningPackageService runningPackageService; private readonly INotificationService notificationService; private readonly ServiceManager vmFactory; @@ -79,12 +79,14 @@ public abstract partial class InferenceGenerationViewModelBase ServiceManager vmFactory, IInferenceClientManager inferenceClientManager, INotificationService notificationService, - ISettingsManager settingsManager + ISettingsManager settingsManager, + RunningPackageService runningPackageService ) : base(notificationService) { this.notificationService = notificationService; this.settingsManager = settingsManager; + this.runningPackageService = runningPackageService; this.vmFactory = vmFactory; ClientManager = inferenceClientManager; @@ -259,6 +261,30 @@ public abstract partial class InferenceGenerationViewModelBase } } + public async Task RunCustomGeneration( + InferenceQueueCustomPromptEventArgs args, + CancellationToken cancellationToken = default + ) + { + if (ClientManager.Client is not { } client) + { + throw new InvalidOperationException("Client is not connected"); + } + + var generationArgs = new ImageGenerationEventArgs + { + Client = client, + Nodes = args.Builder.ToNodeDictionary(), + OutputNodeNames = args.Builder.Connections.OutputNodeNames.ToArray(), + Project = InferenceProjectDocument.FromLoadable(this), + FilesToTransfer = args.FilesToTransfer, + Parameters = new GenerationParameters(), + ClearOutputImages = true + }; + + await RunGeneration(generationArgs, cancellationToken); + } + /// /// Runs a generation task /// @@ -643,12 +669,10 @@ public abstract partial class InferenceGenerationViewModelBase { // Get prompt required extensions // Just static for now but could do manifest lookup when we support custom workflows - var requiredExtensions = nodeDictionary - .ClassTypeRequiredExtensions.Values.SelectMany(x => x) - .ToHashSet(); + var requiredExtensionSpecifiers = nodeDictionary.RequiredExtensions.ToList(); // Skip if no extensions required - if (requiredExtensions.Count == 0) + if (requiredExtensionSpecifiers.Count == 0) { return true; } @@ -661,20 +685,63 @@ public abstract partial class InferenceGenerationViewModelBase await ((GitPackageExtensionManager)manager).GetInstalledExtensionsLiteAsync( localPackagePair.InstalledPackage ) - ).ToImmutableArray(); + ).ToList(); + + var localExtensionsByGitUrl = localExtensions + .Where(ext => ext.GitRepositoryUrl is not null) + .ToDictionary(ext => ext.GitRepositoryUrl!, ext => ext); + + var requiredExtensionReferences = requiredExtensionSpecifiers + .Select(specifier => specifier.Name) + .ToHashSet(); - var missingExtensions = requiredExtensions - .Except(localExtensions.Select(ext => ext.GitRepositoryUrl).WhereNotNull()) - .ToImmutableArray(); + var missingExtensions = new List(); + var outOfDateExtensions = + new List<(ExtensionSpecifier Specifier, InstalledPackageExtension Installed)>(); - if (missingExtensions.Length == 0) + // Check missing extensions and out of date extensions + foreach (var specifier in requiredExtensionSpecifiers) + { + if (!localExtensionsByGitUrl.TryGetValue(specifier.Name, out var localExtension)) + { + missingExtensions.Add(specifier); + continue; + } + + // Check if constraint is specified + if (specifier.Constraint is not null && specifier.TryGetSemVersionRange(out var semVersionRange)) + { + // Get version to compare + localExtension = await manager.GetInstalledExtensionInfoAsync(localExtension); + + // Try to parse local tag to semver + if ( + localExtension.Version?.Tag is not null + && SemVersion.TryParse( + localExtension.Version.Tag, + SemVersionStyles.AllowV, + out var localSemVersion + ) + ) + { + // Check if not satisfied + if (!semVersionRange.Contains(localSemVersion)) + { + outOfDateExtensions.Add((specifier, localExtension)); + } + } + } + } + + if (missingExtensions.Count == 0 && outOfDateExtensions.Count == 0) { return true; } var dialog = DialogHelper.CreateMarkdownDialog( $"#### The following extensions are required for this workflow:\n" - + $"{string.Join("\n- ", missingExtensions)}", + + $"{string.Join("\n- ", missingExtensions.Select(ext => ext.Name))}" + + $"{string.Join("\n- ", outOfDateExtensions.Select(pair => $"{pair.Item1.Name} {pair.Specifier.Constraint} {pair.Specifier.Version} (Current Version: {pair.Installed.Version?.Tag})"))}", "Install Required Extensions?" ); @@ -692,13 +759,14 @@ public abstract partial class InferenceGenerationViewModelBase var steps = new List(); - foreach (var missingExtensionUrl in missingExtensions) + // Add install for missing extensions + foreach (var missingExtension in missingExtensions) { - if (!manifestExtensionsMap.TryGetValue(missingExtensionUrl, out var extension)) + if (!manifestExtensionsMap.TryGetValue(missingExtension.Name, out var extension)) { Logger.Warn( "Extension {MissingExtensionUrl} not found in manifests", - missingExtensionUrl + missingExtension.Name ); continue; } @@ -706,6 +774,18 @@ public abstract partial class InferenceGenerationViewModelBase steps.Add(new InstallExtensionStep(manager, localPackagePair.InstalledPackage, extension)); } + // Add update for out of date extensions + foreach (var (specifier, installed) in outOfDateExtensions) + { + if (!manifestExtensionsMap.TryGetValue(specifier.Name, out var extension)) + { + Logger.Warn("Extension {MissingExtensionUrl} not found in manifests", specifier.Name); + continue; + } + + steps.Add(new UpdateExtensionStep(manager, localPackagePair.InstalledPackage, installed)); + } + var runner = new PackageModificationRunner { ShowDialogOnStart = true, @@ -722,15 +802,12 @@ public abstract partial class InferenceGenerationViewModelBase return; // Restart Package - // TODO: This should be handled by some DI package manager service - var launchPage = App.Services.GetRequiredService(); - try { await Dispatcher.UIThread.InvokeAsync(async () => { - await launchPage.Stop(); - await launchPage.LaunchAsync(); + await runningPackageService.StopPackage(localPackagePair.InstalledPackage.Id); + await runningPackageService.StartPackage(localPackagePair.InstalledPackage); }); } catch (Exception e) diff --git a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceTabViewModelBase.cs b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceTabViewModelBase.cs index 63e5b131..57e6bf47 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Base/InferenceTabViewModelBase.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Base/InferenceTabViewModelBase.cs @@ -98,11 +98,7 @@ public abstract partial class InferenceTabViewModelBase protected async Task SaveViewState() { var eventArgs = new SaveViewStateEventArgs(); - saveViewStateRequestedEventManager?.RaiseEvent( - this, - eventArgs, - nameof(SaveViewStateRequested) - ); + saveViewStateRequestedEventManager?.RaiseEvent(this, eventArgs, nameof(SaveViewStateRequested)); if (eventArgs.StateTask is not { } stateTask) { @@ -128,7 +124,7 @@ public abstract partial class InferenceTabViewModelBase // TODO: Dock reset not working, using this hack for now to get a new view var navService = App.Services.GetRequiredService>(); - navService.NavigateTo(new SuppressNavigationTransitionInfo()); + navService.NavigateTo(new SuppressNavigationTransitionInfo()); ((IPersistentViewProvider)this).AttachedPersistentView = null; navService.NavigateTo(new BetterEntranceNavigationTransition()); } @@ -157,9 +153,7 @@ public abstract partial class InferenceTabViewModelBase if (result == ContentDialogResult.Primary && textFields[0].Text is { } json) { - LoadViewState( - new LoadViewStateEventArgs { State = new ViewState { DockLayout = json } } - ); + LoadViewState(new LoadViewStateEventArgs { State = new ViewState { DockLayout = json } }); } } @@ -226,9 +220,7 @@ public abstract partial class InferenceTabViewModelBase if (this is IParametersLoadableState paramsLoadableVm) { - Dispatcher.UIThread.Invoke( - () => paramsLoadableVm.LoadStateFromParameters(parameters) - ); + Dispatcher.UIThread.Invoke(() => paramsLoadableVm.LoadStateFromParameters(parameters)); } else { diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs index 924d17ed..4754397d 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs @@ -1,13 +1,10 @@ using System; using System.Collections.Generic; using System.Collections.ObjectModel; -using System.ComponentModel; using System.Diagnostics; using System.Linq; using System.Net.Http; -using System.Reactive; using System.Reactive.Linq; -using System.Threading; using System.Threading.Tasks; using AsyncAwaitBestPractices; using Avalonia.Collections; @@ -15,12 +12,15 @@ using Avalonia.Controls; using Avalonia.Controls.Notifications; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; +using DynamicData; +using DynamicData.Alias; +using DynamicData.Binding; using LiteDB; using LiteDB.Async; using NLog; -using OneOf.Types; using Refit; using StabilityMatrix.Avalonia.Languages; +using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.ViewModels.CheckpointManager; @@ -41,24 +41,24 @@ namespace StabilityMatrix.Avalonia.ViewModels.CheckpointBrowser; [View(typeof(CivitAiBrowserPage))] [Singleton] -public partial class CivitAiBrowserViewModel : TabViewModelBase +public partial class CivitAiBrowserViewModel : TabViewModelBase, IInfinitelyScroll { private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); private readonly ICivitApi civitApi; - private readonly IDownloadService downloadService; private readonly ISettingsManager settingsManager; private readonly ServiceManager dialogFactory; private readonly ILiteDbContext liteDbContext; private readonly INotificationService notificationService; private const int MaxModelsPerPage = 20; + private LRUCache< int /* model id */ , CheckpointBrowserCardViewModel - > cache = new(50); + > cache = new(150); [ObservableProperty] - private ObservableCollection? modelCards; + private ObservableCollection modelCards = new(); [ObservableProperty] private DataGridCollectionView? modelCardsView; @@ -81,27 +81,9 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase [ObservableProperty] private CivitModelType selectedModelType = CivitModelType.Checkpoint; - [ObservableProperty] - private int currentPageNumber; - - [ObservableProperty] - private int totalPages; - [ObservableProperty] private bool hasSearched; - [ObservableProperty] - private bool canGoToNextPage; - - [ObservableProperty] - private bool canGoToPreviousPage; - - [ObservableProperty] - private bool canGoToFirstPage; - - [ObservableProperty] - private bool canGoToLastPage; - [ObservableProperty] private bool isIndeterminate; @@ -117,7 +99,8 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase [ObservableProperty] private bool showSantaHats = true; - private List allModelCards = new(); + [ObservableProperty] + private string? nextPageCursor; public IEnumerable AllCivitPeriods => Enum.GetValues(typeof(CivitPeriod)).Cast(); @@ -143,25 +126,11 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase ) { this.civitApi = civitApi; - this.downloadService = downloadService; this.settingsManager = settingsManager; this.dialogFactory = dialogFactory; this.liteDbContext = liteDbContext; this.notificationService = notificationService; - CurrentPageNumber = 1; - CanGoToNextPage = true; - CanGoToLastPage = true; - - Observable - .FromEventPattern(this, nameof(PropertyChanged)) - .Where(x => x.EventArgs.PropertyName == nameof(CurrentPageNumber)) - .Throttle(TimeSpan.FromMilliseconds(250)) - .Select, int>(_ => CurrentPageNumber) - .Where(page => page <= TotalPages && page > 0) - .ObserveOn(SynchronizationContext.Current) - .Subscribe(_ => TrySearchAgain(false).SafeFireAndForget(), err => Logger.Error(err)); - EventManager.Instance.NavigateAndFindCivitModelRequested += OnNavigateAndFindCivitModelRequested; } @@ -171,7 +140,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase return; SearchQuery = $"$#{e}"; - SearchModelsCommand.ExecuteAsync(null).SafeFireAndForget(); + SearchModelsCommand.ExecuteAsync(false).SafeFireAndForget(); } public override void OnLoaded() @@ -223,7 +192,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase /// /// Background update task /// - private async Task CivitModelQuery(CivitModelsRequest request) + private async Task CivitModelQuery(CivitModelsRequest request, bool isInfiniteScroll = false) { var timer = Stopwatch.StartNew(); var queryText = request.Query; @@ -276,15 +245,9 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase } ); - if (cacheNew) - { - Logger.Debug("New cache entry, updating model cards"); - UpdateModelCards(models, modelsResponse.Metadata); - } - else - { - Logger.Debug("Cache entry already exists, not updating model cards"); - } + UpdateModelCards(models, isInfiniteScroll); + + NextPageCursor = modelsResponse.Metadata?.NextCursor; } catch (OperationCanceledException) { @@ -327,7 +290,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase /// /// Updates model cards using api response object. /// - private void UpdateModelCards(IEnumerable? models, CivitMetadata? metadata) + private void UpdateModelCards(List? models, bool addCards = false) { if (models is null) { @@ -335,7 +298,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase } else { - var updateCards = models + var modelsToAdd = models .Select(model => { var cachedViewModel = cache.Get(model.Id); @@ -364,23 +327,34 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase return newCard; }) - .ToList(); + .Where(FilterModelCardsPredicate); - allModelCards = updateCards; - - var filteredCards = updateCards.Where(FilterModelCardsPredicate); if (SortMode == CivitSortMode.Installed) { - filteredCards = filteredCards.OrderByDescending(x => x.UpdateCardText == "Update Available"); + modelsToAdd = modelsToAdd.OrderByDescending(x => x.UpdateCardText == "Update Available"); + } + + if (!addCards) + { + ModelCards = new ObservableCollection(modelsToAdd); } + else + { + foreach (var model in modelsToAdd) + { + if ( + ModelCards.Contains( + model, + new PropertyComparer(x => x.CivitModel.Id) + ) + ) + continue; - ModelCards = new ObservableCollection(filteredCards); + ModelCards.Add(model); + } + } } - TotalPages = metadata?.TotalPages ?? 1; - CanGoToFirstPage = CurrentPageNumber != 1; - CanGoToPreviousPage = CurrentPageNumber > 1; - CanGoToNextPage = CurrentPageNumber < TotalPages; - CanGoToLastPage = CurrentPageNumber != TotalPages; + // Status update ShowMainLoadingSpinner = false; IsIndeterminate = false; @@ -390,27 +364,30 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase private string previousSearchQuery = string.Empty; [RelayCommand] - private async Task SearchModels() + private async Task SearchModels(bool isInfiniteScroll = false) { var timer = Stopwatch.StartNew(); - if (SearchQuery != previousSearchQuery) + if (SearchQuery != previousSearchQuery || !isInfiniteScroll) { // Reset page number - CurrentPageNumber = 1; previousSearchQuery = SearchQuery; + NextPageCursor = null; } // Build request var modelRequest = new CivitModelsRequest { - Limit = MaxModelsPerPage, Nsfw = "true", // Handled by local view filter Sort = SortMode, - Period = SelectedPeriod, - Page = CurrentPageNumber + Period = SelectedPeriod }; + if (NextPageCursor != null) + { + modelRequest.Cursor = NextPageCursor; + } + if (SelectedModelType != CivitModelType.All) { modelRequest.Types = [SelectedModelType]; @@ -516,14 +493,15 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase modelRequest.GetHashCode(), elapsed.TotalSeconds ); - UpdateModelCards(cachedQuery.Items, cachedQuery.Metadata); + NextPageCursor = cachedQuery.Metadata?.NextCursor; + UpdateModelCards(cachedQuery.Items, isInfiniteScroll); // Start remote query (background mode) // Skip when last query was less than 2 min ago var timeSinceCache = DateTimeOffset.UtcNow - cachedQuery.InsertedAt; if (timeSinceCache?.TotalMinutes >= 2) { - CivitModelQuery(modelRequest).SafeFireAndForget(); + CivitModelQuery(modelRequest, isInfiniteScroll).SafeFireAndForget(); Logger.Debug( "Cached query was more than 2 minutes ago ({Seconds:F0} s), updating cache with remote query", timeSinceCache.Value.TotalSeconds @@ -534,54 +512,23 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase { // Not cached, wait for remote query ShowMainLoadingSpinner = true; - await CivitModelQuery(modelRequest); + await CivitModelQuery(modelRequest, isInfiniteScroll); } UpdateResultsText(); } - public void FirstPage() - { - CurrentPageNumber = 1; - } - - public void PreviousPage() - { - if (CurrentPageNumber == 1) - return; - - CurrentPageNumber--; - } - - public void NextPage() - { - if (CurrentPageNumber == TotalPages) - return; - - CurrentPageNumber++; - } - - public void LastPage() - { - CurrentPageNumber = TotalPages; - } - public void ClearSearchQuery() { SearchQuery = string.Empty; } - partial void OnShowNsfwChanged(bool value) + public async Task LoadNextPageAsync() { - settingsManager.Transaction(s => s.ModelBrowserNsfwEnabled, value); - // ModelCardsView?.Refresh(); - var updateCards = allModelCards.Where(FilterModelCardsPredicate); - ModelCards = new ObservableCollection(updateCards); - - if (!HasSearched) - return; - - UpdateResultsText(); + if (NextPageCursor != null) + { + await SearchModelsCommand.ExecuteAsync(true); + } } partial void OnSelectedPeriodChanged(CivitPeriod value) @@ -596,6 +543,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase SelectedBaseModelType ) ); + NextPageCursor = null; } partial void OnSortModeChanged(CivitSortMode value) @@ -610,6 +558,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase SelectedBaseModelType ) ); + NextPageCursor = null; } partial void OnSelectedModelTypeChanged(CivitModelType value) @@ -624,6 +573,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase SelectedBaseModelType ) ); + NextPageCursor = null; } partial void OnSelectedBaseModelTypeChanged(string value) @@ -638,6 +588,7 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase value ) ); + NextPageCursor = null; } private async Task TrySearchAgain(bool shouldUpdatePageNumber = true) @@ -648,18 +599,17 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase if (shouldUpdatePageNumber) { - CurrentPageNumber = 1; + NextPageCursor = null; } // execute command instead of calling method directly so that the IsRunning property gets updated - await SearchModelsCommand.ExecuteAsync(null); + await SearchModelsCommand.ExecuteAsync(false); } private void UpdateResultsText() { NoResultsFound = ModelCards?.Count <= 0; - NoResultsText = - allModelCards.Count > 0 ? $"{allModelCards.Count} results hidden by filters" : "No results found"; + NoResultsText = "No results found"; } public override string Header => Resources.Label_CivitAi; diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs index f3b81788..baebb0fd 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs @@ -8,6 +8,7 @@ using Avalonia.Data; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; using FluentAvalonia.UI.Controls; +using Microsoft.Extensions.DependencyInjection; using NLog; using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.ViewModels.Base; @@ -120,6 +121,7 @@ public partial class CheckpointFile : ViewModelBase if (File.Exists(FilePath)) { IsLoading = true; + Progress = new ProgressReport(0f, "Deleting..."); try { await using var delay = new MinimumDelay(200, 500); @@ -135,6 +137,12 @@ public partial class CheckpointFile : ViewModelBase { await Task.Run(() => File.Delete(cmInfoPath)); } + + var settingsManager = App.Services.GetRequiredService(); + settingsManager.Transaction(s => + { + s.InstalledModelHashes?.Remove(ConnectedModel.Hashes.BLAKE3); + }); } } catch (IOException ex) diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs index ab777a44..e11cccdc 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFolder.cs @@ -111,6 +111,10 @@ public partial class CheckpointFolder : ViewModelBase public IObservableCollection BaseModelOptions { get; } = new ObservableCollectionExtended(); + private IEnumerable allModelOptions = Enum.GetValues() + .Where(x => x != CivitBaseModelType.All) + .Select(x => x.GetStringValue()); + public CheckpointFolder( ISettingsManager settingsManager, IDownloadService downloadService, @@ -175,11 +179,7 @@ public partial class CheckpointFolder : ViewModelBase SubFoldersCache.Refresh(); }); - BaseModelOptionsCache.AddOrUpdate( - Enum.GetValues() - .Where(x => x != CivitBaseModelType.All) - .Select(x => x.GetStringValue()) - ); + BaseModelOptionsCache.AddOrUpdate(allModelOptions); CheckpointFiles.CollectionChanged += OnCheckpointFilesChanged; // DisplayedCheckpointFiles = CheckpointFiles; @@ -187,7 +187,7 @@ public partial class CheckpointFolder : ViewModelBase private bool BaseModelFilter(CheckpointFile file) { - return file.IsConnectedModel + return file.IsConnectedModel && allModelOptions.Contains(file.ConnectedModel!.BaseModel) ? BaseModelOptions.Contains(file.ConnectedModel!.BaseModel) : BaseModelOptions.Contains("Other"); } diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs index f5bc68b9..a3720751 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs @@ -25,6 +25,7 @@ using StabilityMatrix.Core.Models.Api; using StabilityMatrix.Core.Models.Progress; using StabilityMatrix.Core.Processes; using StabilityMatrix.Core.Services; +using Resources = StabilityMatrix.Avalonia.Languages.Resources; using Symbol = FluentIcons.Common.Symbol; using SymbolIconSource = FluentIcons.Avalonia.Fluent.SymbolIconSource; using TeachingTip = StabilityMatrix.Core.Models.Settings.TeachingTip; @@ -43,7 +44,7 @@ public partial class CheckpointsPageViewModel : PageViewModelBase private readonly INotificationService notificationService; private readonly IMetadataImportService metadataImportService; - public override string Title => "Checkpoints"; + public override string Title => Resources.Label_Checkpoints; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Notebook, IsFilled = true }; diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/ConfirmPackageDeleteDialogViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/ConfirmPackageDeleteDialogViewModel.cs new file mode 100644 index 00000000..2d0ce30e --- /dev/null +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/ConfirmPackageDeleteDialogViewModel.cs @@ -0,0 +1,21 @@ +using System; +using CommunityToolkit.Mvvm.ComponentModel; +using StabilityMatrix.Avalonia.ViewModels.Base; +using StabilityMatrix.Avalonia.Views.Dialogs; +using StabilityMatrix.Core.Attributes; + +namespace StabilityMatrix.Avalonia.ViewModels.Dialogs; + +[View(typeof(ConfirmPackageDeleteDialog))] +[ManagedService] +[Transient] +public partial class ConfirmPackageDeleteDialogViewModel : ContentDialogViewModelBase +{ + public required string ExpectedPackageName { get; set; } + + [ObservableProperty] + [NotifyPropertyChangedFor(nameof(IsValid))] + private string packageName = string.Empty; + + public bool IsValid => ExpectedPackageName.Equals(PackageName, StringComparison.Ordinal); +} diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/InferenceConnectionHelpViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/InferenceConnectionHelpViewModel.cs index fafd1206..4f271517 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/InferenceConnectionHelpViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/InferenceConnectionHelpViewModel.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Collections.Immutable; using System.Linq; +using System.Threading.Tasks; using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; @@ -30,6 +31,7 @@ public partial class InferenceConnectionHelpViewModel : ContentDialogViewModelBa private readonly ISettingsManager settingsManager; private readonly INavigationService navigationService; private readonly IPackageFactory packageFactory; + private readonly RunningPackageService runningPackageService; [ObservableProperty] private string title = "Hello"; @@ -58,12 +60,14 @@ public partial class InferenceConnectionHelpViewModel : ContentDialogViewModelBa public InferenceConnectionHelpViewModel( ISettingsManager settingsManager, INavigationService navigationService, - IPackageFactory packageFactory + IPackageFactory packageFactory, + RunningPackageService runningPackageService ) { this.settingsManager = settingsManager; this.navigationService = navigationService; this.packageFactory = packageFactory; + this.runningPackageService = runningPackageService; // Get comfy type installed packages var comfyPackages = this.settingsManager.Settings.InstalledPackages.Where( @@ -122,14 +126,11 @@ public partial class InferenceConnectionHelpViewModel : ContentDialogViewModelBa /// Request launch of the selected package /// [RelayCommand] - private void LaunchSelectedPackage() + private async Task LaunchSelectedPackage() { - if (SelectedPackage?.Id is { } id) + if (SelectedPackage is not null) { - Dispatcher.UIThread.Post(() => - { - EventManager.Instance.OnPackageLaunchRequested(id); - }); + await runningPackageService.StartPackage(SelectedPackage); } } diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/NewOneClickInstallViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/NewOneClickInstallViewModel.cs index 197c17da..f6164bd0 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/NewOneClickInstallViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/NewOneClickInstallViewModel.cs @@ -75,7 +75,7 @@ public partial class NewOneClickInstallViewModel : ContentDialogViewModelBase .Connect() .DeferUntilLoaded() .Filter(incompatiblePredicate) - .Filter(p => p.OfferInOneClickInstaller || ShowIncompatiblePackages) + .Filter(p => p.OfferInOneClickInstaller) .Sort( SortExpressionComparer .Ascending(p => p.InstallerSortOrder) @@ -85,6 +85,10 @@ public partial class NewOneClickInstallViewModel : ContentDialogViewModelBase .Subscribe(); AllPackagesCache.AddOrUpdate(packageFactory.GetAllAvailablePackages()); + if (ShownPackages.Count > 0) + return; + + ShowIncompatiblePackages = true; } [RelayCommand] @@ -182,6 +186,7 @@ public partial class NewOneClickInstallViewModel : ContentDialogViewModelBase { ShowDialogOnStart = false, HideCloseButton = false, + ModificationCompleteMessage = $"{installedPackage.DisplayName} Install Complete", }; runner diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/OneClickInstallViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/OneClickInstallViewModel.cs index 0e28b8fd..d815fb12 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/OneClickInstallViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/OneClickInstallViewModel.cs @@ -192,7 +192,12 @@ public partial class OneClickInstallViewModel : ContentDialogViewModelBase var addInstalledPackageStep = new AddInstalledPackageStep(settingsManager, installedPackage); steps.Add(addInstalledPackageStep); - var runner = new PackageModificationRunner { ShowDialogOnStart = true, HideCloseButton = true, }; + var runner = new PackageModificationRunner + { + ShowDialogOnStart = true, + HideCloseButton = true, + ModificationCompleteMessage = "Install complete" + }; EventManager.Instance.OnPackageInstallProgressAdded(runner); await runner.ExecuteSteps(steps); diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/RecommendedModelsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/RecommendedModelsViewModel.cs index 41750364..e037408e 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/RecommendedModelsViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/RecommendedModelsViewModel.cs @@ -97,9 +97,11 @@ public partial class RecommendedModelsViewModel : ContentDialogViewModelBase new RecommendedModelItemViewModel { ModelVersion = model.ModelVersions.First( - x => !x.BaseModel.Contains("Turbo", StringComparison.OrdinalIgnoreCase) + x => + !x.BaseModel.Contains("Turbo", StringComparison.OrdinalIgnoreCase) + && !x.BaseModel.Contains("Lightning", StringComparison.OrdinalIgnoreCase) ), - Author = $"by {model.Creator.Username}", + Author = $"by {model.Creator?.Username}", CivitModel = model } ) diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs index b992eb6e..9ef4c378 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs @@ -91,7 +91,6 @@ public partial class SelectModelVersionViewModel( { SelectedVersionViewModel = Versions[0]; CanGoToNextImage = true; - LoadInstallLocations(); } partial void OnSelectedVersionViewModelChanged(ModelVersionViewModel? value) @@ -130,6 +129,8 @@ public partial class SelectModelVersionViewModel( var canImport = true; if (settingsManager.IsLibraryDirSet) { + LoadInstallLocations(); + var fileSizeBytes = value?.CivitFile.SizeKb * 1024; var freeSizeBytes = SystemInfo.GetDiskFreeSpaceBytes(settingsManager.ModelsDirectory) ?? long.MaxValue; @@ -151,9 +152,9 @@ public partial class SelectModelVersionViewModel( IsImportEnabled = value?.CivitFile != null && canImport && !ShowEmptyPathWarning; } - partial void OnSelectedInstallLocationChanged(string value) + partial void OnSelectedInstallLocationChanged(string? value) { - if (value.Equals("Custom...", StringComparison.OrdinalIgnoreCase)) + if (value?.Equals("Custom...", StringComparison.OrdinalIgnoreCase) is true) { Dispatcher.UIThread.InvokeAsync(SelectCustomFolder); } @@ -304,14 +305,16 @@ public partial class SelectModelVersionViewModel( private void LoadInstallLocations() { + var installLocations = new ObservableCollection(); + var rootModelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory); var downloadDirectory = rootModelsDirectory.JoinDir( - CivitModel.Type.ConvertTo().GetStringValue() + SelectedFile?.CivitFile.Type == CivitFileType.VAE + ? SharedFolderType.VAE.GetStringValue() + : CivitModel.Type.ConvertTo().GetStringValue() ); - var installLocations = new ObservableCollection - { - downloadDirectory.ToString().Replace(rootModelsDirectory, "Models") - }; + + installLocations.Add(downloadDirectory.ToString().Replace(rootModelsDirectory, "Models")); foreach (var directory in downloadDirectory.EnumerateDirectories()) { diff --git a/StabilityMatrix.Avalonia/ViewModels/FirstLaunchSetupViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/FirstLaunchSetupViewModel.cs index 46815f62..080c29fa 100644 --- a/StabilityMatrix.Avalonia/ViewModels/FirstLaunchSetupViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/FirstLaunchSetupViewModel.cs @@ -3,6 +3,7 @@ using System.Linq; using System.Threading.Tasks; using AsyncAwaitBestPractices; using CommunityToolkit.Mvvm.ComponentModel; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Styles; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.Views; @@ -27,11 +28,9 @@ public partial class FirstLaunchSetupViewModel : ViewModelBase private RefreshBadgeViewModel checkHardwareBadge = new() { - WorkingToolTipText = "We're checking some hardware specifications to determine compatibility.", - SuccessToolTipText = "Everything looks good!", - FailToolTipText = - "We recommend a GPU with CUDA support for the best experience. " - + "You can continue without one, but some packages may not work, and inference may be slower.", + WorkingToolTipText = Resources.Label_CheckingHardware, + SuccessToolTipText = Resources.Label_EverythingLooksGood, + FailToolTipText = Resources.Label_NvidiaGpuRecommended, FailColorBrush = ThemeColors.ThemeYellow, }; diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ControlNetCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ControlNetCardViewModel.cs index 496e90fc..65bab36e 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ControlNetCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ControlNetCardViewModel.cs @@ -1,19 +1,20 @@ using System; using System.ComponentModel.DataAnnotations; using System.Threading.Tasks; +using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; +using DynamicData.Binding; using FluentAvalonia.UI.Controls; using StabilityMatrix.Avalonia.Controls; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.ViewModels.Dialogs; using StabilityMatrix.Core.Attributes; -using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.FileInterfaces; -using StabilityMatrix.Core.Services; +using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; namespace StabilityMatrix.Avalonia.ViewModels.Inference; @@ -32,7 +33,17 @@ public partial class ControlNetCardViewModel : LoadableViewModelBase [ObservableProperty] [Required] - private HybridModelFile? selectedPreprocessor; + private ComfyAuxPreprocessor? selectedPreprocessor; + + [ObservableProperty] + [Required] + [Range(0, 2048)] + private int width; + + [ObservableProperty] + [Required] + [Range(0, 2048)] + private int height; [ObservableProperty] [Required] @@ -62,6 +73,18 @@ public partial class ControlNetCardViewModel : LoadableViewModelBase ClientManager = clientManager; SelectImageCardViewModel = vmFactory.Get(); + + // Update our width and height when the image changes + SelectImageCardViewModel + .WhenPropertyChanged(card => card.CurrentBitmapSize) + .Subscribe(propertyValue => + { + if (!propertyValue.Value.IsEmpty) + { + Width = propertyValue.Value.Width; + Height = propertyValue.Value.Height; + } + }); } [RelayCommand] @@ -79,4 +102,57 @@ public partial class ControlNetCardViewModel : LoadableViewModelBase confirmDialog.StartDownload(); } } + + [RelayCommand] + private async Task PreviewPreprocessor(ComfyAuxPreprocessor? preprocessor) + { + if ( + preprocessor is null + || SelectImageCardViewModel.ImageSource is not { } imageSource + || SelectImageCardViewModel.IsImageFileNotFound + ) + return; + + var args = new InferenceQueueCustomPromptEventArgs(); + + var images = SelectImageCardViewModel.GetInputImages(); + + await ClientManager.UploadInputImageAsync(imageSource); + + var image = args.Nodes.AddTypedNode( + new ComfyNodeBuilder.LoadImage + { + Name = args.Nodes.GetUniqueName("Preprocessor_LoadImage"), + Image = + SelectImageCardViewModel.ImageSource?.GetHashGuidFileNameCached("Inference") + ?? throw new ValidationException("No ImageSource") + } + ).Output1; + + var aioPreprocessor = args.Nodes.AddTypedNode( + new ComfyNodeBuilder.AIOPreprocessor + { + Name = args.Nodes.GetUniqueName("Preprocessor"), + Image = image, + Preprocessor = preprocessor.ToString(), + Resolution = Width is <= 2048 and > 0 ? Width : 512 + } + ); + + args.Builder.Connections.OutputNodes.Add( + args.Nodes.AddTypedNode( + new ComfyNodeBuilder.PreviewImage + { + Name = args.Nodes.GetUniqueName("Preprocessor_OutputImage"), + Images = aioPreprocessor.Output + } + ) + ); + + // Queue + Dispatcher.UIThread.Post(() => EventManager.Instance.OnInferenceQueueCustomPrompt(args)); + + // We don't know when it's done so wait a bit? + await Task.Delay(1000); + } } diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs index 50c4734f..77c5230f 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs @@ -27,9 +27,17 @@ public class InferenceImageToImageViewModel : InferenceTextToImageViewModel IInferenceClientManager inferenceClientManager, INotificationService notificationService, ISettingsManager settingsManager, - IModelIndexService modelIndexService + IModelIndexService modelIndexService, + RunningPackageService runningPackageService ) - : base(notificationService, inferenceClientManager, settingsManager, vmFactory, modelIndexService) + : base( + notificationService, + inferenceClientManager, + settingsManager, + vmFactory, + modelIndexService, + runningPackageService + ) { SelectImageCardViewModel = vmFactory.Get(); @@ -77,12 +85,12 @@ public class InferenceImageToImageViewModel : InferenceTextToImageViewModel var mainImages = SelectImageCardViewModel.GetInputImages(); var samplerImages = SamplerCardViewModel - .ModulesCardViewModel - .Cards - .OfType() + .ModulesCardViewModel.Cards.OfType() .SelectMany(m => m.GetInputImages()); - var moduleImages = ModulesCardViewModel.Cards.OfType().SelectMany(m => m.GetInputImages()); + var moduleImages = ModulesCardViewModel + .Cards.OfType() + .SelectMany(m => m.GetInputImages()); return mainImages.Concat(samplerImages).Concat(moduleImages); } diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs index b57e3406..035774c1 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToVideoViewModel.cs @@ -65,9 +65,10 @@ public partial class InferenceImageToVideoViewModel IInferenceClientManager inferenceClientManager, ISettingsManager settingsManager, ServiceManager vmFactory, - IModelIndexService modelIndexService + IModelIndexService modelIndexService, + RunningPackageService runningPackageService ) - : base(vmFactory, inferenceClientManager, notificationService, settingsManager) + : base(vmFactory, inferenceClientManager, notificationService, settingsManager, runningPackageService) { this.notificationService = notificationService; this.modelIndexService = modelIndexService; diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageUpscaleViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageUpscaleViewModel.cs index 2ca1dbd8..540de0d9 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageUpscaleViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageUpscaleViewModel.cs @@ -59,9 +59,10 @@ public class InferenceImageUpscaleViewModel : InferenceGenerationViewModelBase INotificationService notificationService, IInferenceClientManager inferenceClientManager, ISettingsManager settingsManager, - ServiceManager vmFactory + ServiceManager vmFactory, + RunningPackageService runningPackageService ) - : base(vmFactory, inferenceClientManager, notificationService, settingsManager) + : base(vmFactory, inferenceClientManager, notificationService, settingsManager, runningPackageService) { this.notificationService = notificationService; @@ -142,7 +143,10 @@ public class InferenceImageUpscaleViewModel : InferenceGenerationViewModelBase } /// - protected override async Task GenerateImageImpl(GenerateOverrides overrides, CancellationToken cancellationToken) + protected override async Task GenerateImageImpl( + GenerateOverrides overrides, + CancellationToken cancellationToken + ) { if (!ClientManager.IsConnected) { @@ -169,7 +173,10 @@ public class InferenceImageUpscaleViewModel : InferenceGenerationViewModelBase Client = ClientManager.Client, Nodes = buildPromptArgs.Builder.ToNodeDictionary(), OutputNodeNames = buildPromptArgs.Builder.Connections.OutputNodeNames.ToArray(), - Parameters = new GenerationParameters { ModelName = UpscalerCardViewModel.SelectedUpscaler?.Name, }, + Parameters = new GenerationParameters + { + ModelName = UpscalerCardViewModel.SelectedUpscaler?.Name, + }, Project = InferenceProjectDocument.FromLoadable(this) }; diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs index 12d0d5fc..a1924559 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs @@ -59,9 +59,10 @@ public class InferenceTextToImageViewModel : InferenceGenerationViewModelBase, I IInferenceClientManager inferenceClientManager, ISettingsManager settingsManager, ServiceManager vmFactory, - IModelIndexService modelIndexService + IModelIndexService modelIndexService, + RunningPackageService runningPackageService ) - : base(vmFactory, inferenceClientManager, notificationService, settingsManager) + : base(vmFactory, inferenceClientManager, notificationService, settingsManager, runningPackageService) { this.notificationService = notificationService; this.modelIndexService = modelIndexService; diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/ControlNetModule.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/ControlNetModule.cs index 23260daf..77955192 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/ControlNetModule.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/ControlNetModule.cs @@ -6,6 +6,9 @@ using StabilityMatrix.Avalonia.Models.Inference; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Core.Attributes; +using StabilityMatrix.Core.Extensions; +using StabilityMatrix.Core.Helper; +using StabilityMatrix.Core.Models.Api.Comfy; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; namespace StabilityMatrix.Avalonia.ViewModels.Inference.Modules; @@ -39,7 +42,7 @@ public class ControlNetModule : ModuleBase { var card = GetCard(); - var imageLoad = e.Nodes.AddTypedNode( + var image = e.Nodes.AddTypedNode( new ComfyNodeBuilder.LoadImage { Name = e.Nodes.GetUniqueName("ControlNet_LoadImage"), @@ -47,7 +50,101 @@ public class ControlNetModule : ModuleBase card.SelectImageCardViewModel.ImageSource?.GetHashGuidFileNameCached("Inference") ?? throw new ValidationException("No ImageSource") } - ); + ).Output1; + + if (card.SelectedPreprocessor is { } preprocessor && preprocessor != ComfyAuxPreprocessor.None) + { + var aioPreprocessor = e.Nodes.AddTypedNode( + new ComfyNodeBuilder.AIOPreprocessor + { + Name = e.Nodes.GetUniqueName("ControlNet_Preprocessor"), + Image = image, + Preprocessor = preprocessor.ToString(), + // Use width if valid, else default of 512 + Resolution = card.Width is <= 2048 and > 0 ? card.Width : 512 + } + ); + + image = aioPreprocessor.Output; + } + + // If ReferenceOnly is selected, use special node + if (card.SelectedModel == RemoteModels.ControlNetReferenceOnlyModel) + { + // We need to rescale image to be the current primary size if it's not already + var originalPrimary = e.Temp.Primary!.Unwrap(); + var originalPrimarySize = e.Builder.Connections.PrimarySize; + + if (card.SelectImageCardViewModel.CurrentBitmapSize != originalPrimarySize) + { + var scaled = e.Builder.Group_Upscale( + e.Nodes.GetUniqueName("ControlNet_Rescale"), + image, + e.Temp.GetDefaultVAE(), + ComfyUpscaler.NearestExact, + originalPrimarySize.Width, + originalPrimarySize.Height + ); + e.Temp.Primary = scaled; + } + else + { + e.Temp.Primary = image; + } + + // Set image as new latent source, add reference only node + var model = e.Temp.GetRefinerOrBaseModel(); + var controlNetReferenceOnly = e.Nodes.AddTypedNode( + new ComfyNodeBuilder.ReferenceOnlySimple + { + Name = e.Nodes.GetUniqueName("ControlNet_ReferenceOnly"), + Reference = e.Builder.GetPrimaryAsLatent( + e.Temp.Primary, + e.Builder.Connections.GetDefaultVAE() + ), + Model = model + } + ); + + var referenceOnlyModel = controlNetReferenceOnly.Output1; + + // If ControlNet strength is not 1, add Model Merge + if (Math.Abs(card.Strength - 1) > 0.01) + { + var modelBlend = e.Nodes.AddTypedNode( + new ComfyNodeBuilder.ModelMergeSimple + { + Name = e.Nodes.GetUniqueName("ControlNet_ReferenceOnly_ModelMerge"), + Model1 = referenceOnlyModel, + Model2 = e.Temp.GetRefinerOrBaseModel(), + // Where 0 is full reference only, 1 is full original + Ratio = 1 - card.Strength + } + ); + + referenceOnlyModel = modelBlend.Output; + } + + // Set output as new primary and model source + if (model == e.Temp.Refiner.Model) + { + e.Temp.Refiner.Model = referenceOnlyModel; + } + else + { + e.Temp.Base.Model = referenceOnlyModel; + } + e.Temp.Primary = controlNetReferenceOnly.Output2; + + // Indicate that the Primary latent has been temp batched + // https://github.com/comfyanonymous/ComfyUI_experiments/issues/11 + + e.Temp.IsPrimaryTempBatched = true; + // Index 0 is the original image, index 1 is the reference only latent + e.Temp.PrimaryTempBatchPickIndex = 1; + + return; + } var controlNetLoader = e.Nodes.AddTypedNode( new ComfyNodeBuilder.ControlNetLoader @@ -62,36 +159,36 @@ public class ControlNetModule : ModuleBase new ComfyNodeBuilder.ControlNetApplyAdvanced { Name = e.Nodes.GetUniqueName("ControlNetApply"), - Image = imageLoad.Output1, + Image = image, ControlNet = controlNetLoader.Output, - Positive = e.Temp.Conditioning?.Positive ?? throw new ArgumentException("No Conditioning"), - Negative = e.Temp.Conditioning?.Negative ?? throw new ArgumentException("No Conditioning"), + Positive = e.Temp.Base.Conditioning!.Unwrap().Positive, + Negative = e.Temp.Base.Conditioning.Negative, Strength = card.Strength, StartPercent = card.StartPercent, EndPercent = card.EndPercent, } ); - e.Temp.Conditioning = (controlNetApply.Output1, controlNetApply.Output2); + e.Temp.Base.Conditioning = (controlNetApply.Output1, controlNetApply.Output2); // Refiner if available - if (e.Temp.RefinerConditioning is not null) + if (e.Temp.Refiner.Conditioning is not null) { var controlNetRefinerApply = e.Nodes.AddTypedNode( new ComfyNodeBuilder.ControlNetApplyAdvanced { Name = e.Nodes.GetUniqueName("Refiner_ControlNetApply"), - Image = imageLoad.Output1, + Image = image, ControlNet = controlNetLoader.Output, - Positive = e.Temp.RefinerConditioning.Positive, - Negative = e.Temp.RefinerConditioning.Negative, + Positive = e.Temp.Refiner.Conditioning!.Unwrap().Positive, + Negative = e.Temp.Refiner.Conditioning.Negative, Strength = card.Strength, StartPercent = card.StartPercent, EndPercent = card.EndPercent, } ); - e.Temp.RefinerConditioning = (controlNetRefinerApply.Output1, controlNetRefinerApply.Output2); + e.Temp.Refiner.Conditioning = (controlNetRefinerApply.Output1, controlNetRefinerApply.Output2); } } } diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/HiresFixModule.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/HiresFixModule.cs index 588cd263..94a0d1bc 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/HiresFixModule.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/Modules/HiresFixModule.cs @@ -78,30 +78,39 @@ public partial class HiresFixModule : ModuleBase ); } - var hiresSampler = builder - .Nodes - .AddTypedNode( - new ComfyNodeBuilder.KSampler - { - Name = builder.Nodes.GetUniqueName("HiresFix_Sampler"), - Model = builder.Connections.GetRefinerOrBaseModel(), - Seed = builder.Connections.Seed, - Steps = samplerCard.Steps, - Cfg = samplerCard.CfgScale, - SamplerName = - samplerCard.SelectedSampler?.Name - ?? e.Builder.Connections.PrimarySampler?.Name - ?? throw new ArgumentException("No PrimarySampler"), - Scheduler = - samplerCard.SelectedScheduler?.Name - ?? e.Builder.Connections.PrimaryScheduler?.Name - ?? throw new ArgumentException("No PrimaryScheduler"), - Positive = builder.Connections.GetRefinerOrBaseConditioning().Positive, - Negative = builder.Connections.GetRefinerOrBaseConditioning().Negative, - LatentImage = builder.GetPrimaryAsLatent(), - Denoise = samplerCard.DenoiseStrength - } - ); + // If we need to inherit primary sampler addons, use their temp args + if (samplerCard.InheritPrimarySamplerAddons) + { + e.Temp = e.Builder.Connections.BaseSamplerTemporaryArgs ?? e.CreateTempFromBuilder(); + } + else + { + // otherwise just use new ones + e.Temp = e.CreateTempFromBuilder(); + } + + var hiresSampler = builder.Nodes.AddTypedNode( + new ComfyNodeBuilder.KSampler + { + Name = builder.Nodes.GetUniqueName("HiresFix_Sampler"), + Model = builder.Connections.GetRefinerOrBaseModel(), + Seed = builder.Connections.Seed, + Steps = samplerCard.Steps, + Cfg = samplerCard.CfgScale, + SamplerName = + samplerCard.SelectedSampler?.Name + ?? e.Builder.Connections.PrimarySampler?.Name + ?? throw new ArgumentException("No PrimarySampler"), + Scheduler = + samplerCard.SelectedScheduler?.Name + ?? e.Builder.Connections.PrimaryScheduler?.Name + ?? throw new ArgumentException("No PrimaryScheduler"), + Positive = e.Temp.GetRefinerOrBaseConditioning().Positive, + Negative = e.Temp.GetRefinerOrBaseConditioning().Negative, + LatentImage = builder.GetPrimaryAsLatent(), + Denoise = samplerCard.DenoiseStrength + } + ); // Set as primary builder.Connections.Primary = hiresSampler.Output; diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs index c2409259..8890a514 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs @@ -87,6 +87,11 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo [Required] private ComfyScheduler? selectedScheduler = ComfyScheduler.Normal; + [ObservableProperty] + [property: Category("Settings")] + [property: DisplayName("Inherit Primary Sampler Addons")] + private bool inheritPrimarySamplerAddons = true; + [JsonPropertyName("Modules")] public StackEditableCardViewModel ModulesCardViewModel { get; } @@ -130,8 +135,7 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo } // Provide temp values - e.Temp.Conditioning = e.Builder.Connections.Base.Conditioning; - e.Temp.RefinerConditioning = e.Builder.Connections.Refiner.Conditioning; + e.Temp = e.CreateTempFromBuilder(); // Apply steps from our addons ApplyAddonSteps(e); @@ -142,17 +146,26 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo if (!e.Nodes.ContainsKey("Sampler")) { ApplyStepsInitialSampler(e); + + // Save temp + e.Builder.Connections.BaseSamplerTemporaryArgs = e.Temp; } else { - ApplyStepsAdditionalSampler(e); + // Hires does its own sampling so just throw I guess + throw new InvalidOperationException( + "Sampler ApplyStep was called when Sampler node already exists" + ); } } private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e) { // Get primary as latent using vae - var primaryLatent = e.Builder.GetPrimaryAsLatent(); + var primaryLatent = e.Builder.GetPrimaryAsLatent( + e.Temp.Primary!.Unwrap(), + e.Builder.Connections.GetDefaultVAE() + ); // Set primary sampler and scheduler var primarySampler = SelectedSampler ?? throw new ValidationException("Sampler not selected"); @@ -162,8 +175,8 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo e.Builder.Connections.PrimaryScheduler = primaryScheduler; // Use Temp Conditioning that may be modified by addons - var conditioning = e.Temp.Conditioning.Unwrap(); - var refinerConditioning = e.Temp.RefinerConditioning; + var conditioning = e.Temp.Base.Conditioning.Unwrap(); + var refinerConditioning = e.Temp.Refiner.Conditioning; // Use custom sampler if SDTurbo scheduler is selected if (e.Builder.Connections.PrimaryScheduler == ComfyScheduler.SDTurbo) @@ -209,21 +222,16 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo ); e.Builder.Connections.Primary = sampler.Output1; - - return; } - // Use KSampler if no refiner, otherwise need KSamplerAdvanced - if (e.Builder.Connections.Refiner.Model is null) + else if (e.Builder.Connections.Refiner.Model is null) { - var baseConditioning = e.Builder.Connections.Base.Conditioning.Unwrap(); - // No refiner var sampler = e.Nodes.AddTypedNode( new ComfyNodeBuilder.KSampler { Name = "Sampler", - Model = e.Builder.Connections.Base.Model.Unwrap(), + Model = e.Temp.Base.Model!.Unwrap(), Seed = e.Builder.Connections.Seed, SamplerName = primarySampler.Name, Scheduler = primaryScheduler.Name, @@ -245,7 +253,7 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo new ComfyNodeBuilder.KSamplerAdvanced { Name = "Sampler", - Model = e.Builder.Connections.Base.Model.Unwrap(), + Model = e.Temp.Base.Model!.Unwrap(), AddNoise = true, NoiseSeed = e.Builder.Connections.Seed, Steps = TotalSteps, @@ -261,8 +269,30 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo } ); + e.Builder.Connections.Primary = sampler.Output; + } + + // If temp batched, add a LatentFromBatch to pick the temp batch right after first sampler + if (e.Temp.IsPrimaryTempBatched) + { + e.Builder.Connections.Primary = e.Nodes.AddTypedNode( + new ComfyNodeBuilder.LatentFromBatch + { + Name = e.Nodes.GetUniqueName("ControlNet_LatentFromBatch"), + Samples = e.Builder.GetPrimaryAsLatent(), + BatchIndex = e.Temp.PrimaryTempBatchPickIndex, + // Use max length here as recommended + // https://github.com/comfyanonymous/ComfyUI_experiments/issues/11 + Length = 64 + } + ).Output; + } + + // Refiner + if (e.Builder.Connections.Refiner.Model is not null) + { // Add refiner sampler - var refinerSampler = e.Nodes.AddTypedNode( + e.Builder.Connections.Primary = e.Nodes.AddTypedNode( new ComfyNodeBuilder.KSamplerAdvanced { Name = "Sampler_Refiner", @@ -276,19 +306,15 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo Positive = refinerConditioning!.Positive, Negative = refinerConditioning.Negative, // Connect to previous sampler - LatentImage = sampler.Output, + LatentImage = e.Builder.GetPrimaryAsLatent(), StartAtStep = Steps, EndAtStep = TotalSteps, ReturnWithLeftoverNoise = false } - ); - - e.Builder.Connections.Primary = refinerSampler.Output; + ).Output; } } - private void ApplyStepsAdditionalSampler(ModuleApplyStepEventArgs e) { } - /// /// Applies each step of our addons /// diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs index 6e2e7b15..98c637e5 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs @@ -90,9 +90,18 @@ public partial class SelectImageCardViewModel(INotificationService notificationS partial void OnImageSourceChanged(ImageSource? value) { // Cache the hash for later upload use - if (value?.LocalFile is { Exists: true }) + if (value?.LocalFile is { Exists: true } localFile) { - value.GetBlake3HashAsync().SafeFireAndForget(); + value + .GetBlake3HashAsync() + .SafeFireAndForget(ex => + { + Logger.Warn(ex, "Error getting hash for image {Path}", localFile.Name); + notificationService.ShowPersistent( + $"Error getting hash for image {localFile.Name}", + $"{ex.GetType().Name}: {ex.Message}" + ); + }); } } diff --git a/StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs index c36fa128..c641f18f 100644 --- a/StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/InferenceViewModel.cs @@ -1,6 +1,8 @@ using System; +using System.Collections.Generic; using System.Collections.Immutable; using System.Collections.ObjectModel; +using System.Collections.Specialized; using System.Linq; using System.Reactive.Linq; using System.Text.Json; @@ -16,6 +18,7 @@ using CommunityToolkit.Mvvm.Input; using FluentAvalonia.UI.Controls; using NLog; using StabilityMatrix.Avalonia.Extensions; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; @@ -51,8 +54,10 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable private readonly ServiceManager vmFactory; private readonly IModelIndexService modelIndexService; private readonly ILiteDbContext liteDbContext; + private readonly RunningPackageService runningPackageService; + private Guid? selectedPackageId; - public override string Title => "Inference"; + public override string Title => Resources.Label_Inference; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.AppGeneric, IsFilled = true }; @@ -62,7 +67,7 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable State = ProgressState.Failed, FailToolTipText = "Not connected", FailIcon = FluentAvalonia.UI.Controls.Symbol.Refresh, - SuccessToolTipText = "Connected", + SuccessToolTipText = Resources.Label_Connected, }; public IInferenceClientManager ClientManager { get; } @@ -86,6 +91,8 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable public bool IsComfyRunning => RunningPackage?.BasePackage is ComfyUI; + private IDisposable? onStartupComplete; + public InferenceViewModel( ServiceManager vmFactory, INotificationService notificationService, @@ -93,6 +100,7 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable ISettingsManager settingsManager, IModelIndexService modelIndexService, ILiteDbContext liteDbContext, + RunningPackageService runningPackageService, SharedState sharedState ) { @@ -101,12 +109,13 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable this.settingsManager = settingsManager; this.modelIndexService = modelIndexService; this.liteDbContext = liteDbContext; + this.runningPackageService = runningPackageService; ClientManager = inferenceClientManager; SharedState = sharedState; // Keep RunningPackage updated with the current package pair - EventManager.Instance.RunningPackageStatusChanged += OnRunningPackageStatusChanged; + runningPackageService.RunningPackages.CollectionChanged += RunningPackagesOnCollectionChanged; // "Send to Inference" EventManager.Instance.InferenceTextToImageRequested += OnInferenceTextToImageRequested; @@ -114,58 +123,112 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable EventManager.Instance.InferenceImageToImageRequested += OnInferenceImageToImageRequested; EventManager.Instance.InferenceImageToVideoRequested += OnInferenceImageToVideoRequested; + // Global requests for custom prompt queueing + EventManager.Instance.InferenceQueueCustomPrompt += OnInferenceQueueCustomPromptRequested; + MenuSaveAsCommand.WithConditionalNotificationErrorHandler(notificationService); MenuOpenProjectCommand.WithConditionalNotificationErrorHandler(notificationService); } + private void DisconnectFromComfy() + { + RunningPackage = null; + + // Cancel any pending connection + if (ConnectCancelCommand.CanExecute(null)) + { + ConnectCancelCommand.Execute(null); + } + onStartupComplete?.Dispose(); + onStartupComplete = null; + IsWaitingForConnection = false; + + // Disconnect + Logger.Trace("On package close - disconnecting"); + DisconnectCommand.Execute(null); + } + /// /// Updates the RunningPackage property when the running package changes. /// Also starts a connection to the backend if a new ComfyUI package is running. /// And disconnects if the package is closed. /// - private void OnRunningPackageStatusChanged(object? sender, RunningPackageStatusChangedEventArgs e) + private void RunningPackagesOnCollectionChanged(object? sender, NotifyCollectionChangedEventArgs e) { - RunningPackage = e.CurrentPackagePair; + if ( + e.NewItems?.OfType>().Select(x => x.Value) + is not { } newItems + ) + { + if (RunningPackage != null) + { + DisconnectFromComfy(); + } + return; + } - IDisposable? onStartupComplete = null; + var comfyViewModel = newItems.FirstOrDefault( + vm => + vm.RunningPackage.InstalledPackage.Id == selectedPackageId + || vm.RunningPackage.BasePackage is ComfyUI + ); - Dispatcher.UIThread.Post(() => + if (comfyViewModel is null && RunningPackage?.BasePackage is ComfyUI) { - if (e.CurrentPackagePair?.BasePackage is ComfyUI package) - { - IsWaitingForConnection = true; - onStartupComplete = Observable - .FromEventPattern(package, nameof(package.StartupComplete)) - .Take(1) - .Subscribe(_ => + DisconnectFromComfy(); + } + else if (comfyViewModel != null && RunningPackage == null) + { + IsWaitingForConnection = true; + RunningPackage = comfyViewModel.RunningPackage; + onStartupComplete = Observable + .FromEventPattern( + comfyViewModel.RunningPackage.BasePackage, + nameof(comfyViewModel.RunningPackage.BasePackage.StartupComplete) + ) + .Take(1) + .Subscribe(_ => + { + Dispatcher.UIThread.Post(() => { - Dispatcher.UIThread.Post(() => + if (ConnectCommand.CanExecute(null)) { - if (ConnectCommand.CanExecute(null)) - { - Logger.Trace("On package launch - starting connection"); - ConnectCommand.Execute(null); - } - IsWaitingForConnection = false; - }); + Logger.Trace("On package launch - starting connection"); + ConnectCommand.Execute(null); + } + + IsWaitingForConnection = false; }); - } - else - { - // Cancel any pending connection - if (ConnectCancelCommand.CanExecute(null)) + }); + } + } + + private void OnInferenceQueueCustomPromptRequested(object? sender, InferenceQueueCustomPromptEventArgs e) + { + // Get currently selected tab + var currentTab = SelectedTab; + + if (currentTab is InferenceGenerationViewModelBase generationViewModel) + { + Dispatcher + .UIThread.InvokeAsync(async () => { - ConnectCancelCommand.Execute(null); - } - onStartupComplete?.Dispose(); - onStartupComplete = null; - IsWaitingForConnection = false; + await generationViewModel.RunCustomGeneration(e); + }) + .SafeFireAndForget(ex => + { + Logger.Error(ex, "Failed to queue prompt"); - // Disconnect - Logger.Trace("On package close - disconnecting"); - DisconnectCommand.Execute(null); - } - }); + Dispatcher.UIThread.Post(() => + { + notificationService.ShowPersistent( + "Failed to queue prompt", + $"{ex.GetType().Name}: {ex.Message}", + NotificationType.Error + ); + }); + }); + } } public override void OnLoaded() @@ -390,7 +453,12 @@ public partial class InferenceViewModel : PageViewModelBase, IAsyncDisposable private async Task ShowConnectionHelp() { var vm = vmFactory.Get(); - await vm.CreateDialog().ShowAsync(); + var result = await vm.CreateDialog().ShowAsync(); + + if (result != ContentDialogResult.Primary) + return; + + selectedPackageId = vm.SelectedPackage?.Id; } /// diff --git a/StabilityMatrix.Avalonia/ViewModels/MainWindowViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/MainWindowViewModel.cs index 5aec5637..2b77217a 100644 --- a/StabilityMatrix.Avalonia/ViewModels/MainWindowViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/MainWindowViewModel.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; +using System.IO; using System.Linq; using System.Threading.Tasks; using AsyncAwaitBestPractices; @@ -36,6 +37,7 @@ public partial class MainWindowViewModel : ViewModelBase private readonly ITrackedDownloadService trackedDownloadService; private readonly IDiscordRichPresenceService discordRichPresenceService; private readonly IModelIndexService modelIndexService; + private readonly Lazy modelDownloadLinkHandler; public string Greeting => "Welcome to Avalonia!"; [ObservableProperty] @@ -71,7 +73,8 @@ public partial class MainWindowViewModel : ViewModelBase IDiscordRichPresenceService discordRichPresenceService, ServiceManager dialogFactory, ITrackedDownloadService trackedDownloadService, - IModelIndexService modelIndexService + IModelIndexService modelIndexService, + Lazy modelDownloadLinkHandler ) { this.settingsManager = settingsManager; @@ -79,6 +82,7 @@ public partial class MainWindowViewModel : ViewModelBase this.discordRichPresenceService = discordRichPresenceService; this.trackedDownloadService = trackedDownloadService; this.modelIndexService = modelIndexService; + this.modelDownloadLinkHandler = modelDownloadLinkHandler; ProgressManagerViewModel = dialogFactory.Get(); UpdateViewModel = dialogFactory.Get(); } @@ -107,6 +111,25 @@ public partial class MainWindowViewModel : ViewModelBase return; } + try + { + await modelDownloadLinkHandler.Value.StartListening(); + } + catch (IOException) + { + var dialog = new BetterContentDialog + { + Title = Resources.Label_StabilityMatrixAlreadyRunning, + Content = Resources.Label_AnotherInstanceAlreadyRunning, + IsPrimaryButtonEnabled = true, + PrimaryButtonText = Resources.Action_Close, + DefaultButton = ContentDialogButton.Primary + }; + await dialog.ShowAsync(); + App.Shutdown(); + return; + } + // Initialize Discord Rich Presence (this needs LibraryDir so is set here) discordRichPresenceService.UpdateState(); @@ -139,14 +162,6 @@ public partial class MainWindowViewModel : ViewModelBase Content = new NewOneClickInstallDialog { DataContext = viewModel }, }; - EventManager.Instance.OneClickInstallFinished += (_, skipped) => - { - if (skipped) - return; - - EventManager.Instance.OnTeachingTooltipNeeded(); - }; - var firstDialogResult = await dialog.ShowAsync(App.TopLevel); if (firstDialogResult != ContentDialogResult.Primary) diff --git a/StabilityMatrix.Avalonia/ViewModels/NewCheckpointsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/NewCheckpointsPageViewModel.cs index 45ea5ed0..ba95645c 100644 --- a/StabilityMatrix.Avalonia/ViewModels/NewCheckpointsPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/NewCheckpointsPageViewModel.cs @@ -14,6 +14,7 @@ using FluentAvalonia.UI.Controls; using Microsoft.Extensions.Logging; using Refit; using StabilityMatrix.Avalonia.Controls; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.ViewModels.CheckpointManager; @@ -47,7 +48,7 @@ public partial class NewCheckpointsPageViewModel( IMetadataImportService metadataImportService ) : PageViewModelBase { - public override string Title => "Checkpoint Manager"; + public override string Title => Resources.Label_CheckpointManager; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Cellular5g, IsFilled = true }; diff --git a/StabilityMatrix.Avalonia/ViewModels/NewPackageManagerViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/NewPackageManagerViewModel.cs index d900f2bc..d96dac06 100644 --- a/StabilityMatrix.Avalonia/ViewModels/NewPackageManagerViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/NewPackageManagerViewModel.cs @@ -3,6 +3,7 @@ using System.Collections.ObjectModel; using CommunityToolkit.Mvvm.ComponentModel; using DynamicData; using FluentAvalonia.UI.Controls; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.ViewModels.Dialogs; @@ -18,7 +19,7 @@ namespace StabilityMatrix.Avalonia.ViewModels; [Singleton] public partial class NewPackageManagerViewModel : PageViewModelBase { - public override string Title => "Packages"; + public override string Title => Resources.Label_Packages; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Box, IsFilled = true }; public IReadOnlyList SubPages { get; } @@ -58,6 +59,10 @@ public partial class NewPackageManagerViewModel : PageViewModelBase { CurrentPagePath.Add(value); } + else if (value is RunningPackageViewModel) + { + CurrentPagePath.Add(value); + } else { CurrentPagePath.Clear(); diff --git a/StabilityMatrix.Avalonia/ViewModels/OutputsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/OutputsPageViewModel.cs index ecb81a55..8dcecc60 100644 --- a/StabilityMatrix.Avalonia/ViewModels/OutputsPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/OutputsPageViewModel.cs @@ -5,6 +5,7 @@ using System.Diagnostics; using System.IO; using System.Linq; using System.Reactive.Linq; +using System.Threading; using System.Threading.Tasks; using AsyncAwaitBestPractices; using AsyncImageLoader; @@ -28,8 +29,6 @@ using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Avalonia.ViewModels.Dialogs; using StabilityMatrix.Avalonia.ViewModels.OutputsPage; using StabilityMatrix.Core.Attributes; -using StabilityMatrix.Core.Exceptions; -using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Helper.Factory; using StabilityMatrix.Core.Models; @@ -52,19 +51,21 @@ public partial class OutputsPageViewModel : PageViewModelBase private readonly INotificationService notificationService; private readonly INavigationService navigationService; private readonly ILogger logger; + private readonly List cancellationTokenSources = []; + public override string Title => Resources.Label_OutputsPageTitle; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Grid, IsFilled = true }; public SourceCache OutputsCache { get; } = new(file => file.AbsolutePath); + private SourceCache categoriesCache = new(category => category.Path); + public IObservableCollection Outputs { get; set; } = new ObservableCollectionExtended(); - public IEnumerable OutputTypes { get; } = Enum.GetValues(); - - [ObservableProperty] - private ObservableCollection categories; + public IObservableCollection Categories { get; set; } = + new ObservableCollectionExtended(); [ObservableProperty] [NotifyPropertyChangedFor(nameof(CanShowOutputTypes))] @@ -86,6 +87,15 @@ public partial class OutputsPageViewModel : PageViewModelBase [ObservableProperty] private bool isConsolidating; + [ObservableProperty] + private bool isLoading; + + [ObservableProperty] + private bool showFolders; + + [ObservableProperty] + private bool isChangingCategory; + public bool CanShowOutputTypes => SelectedCategory?.Name?.Equals("Shared Output Folder") ?? false; public string NumImagesSelected => @@ -95,6 +105,8 @@ public partial class OutputsPageViewModel : PageViewModelBase private string[] allowedExtensions = [".png", ".webp"]; + private PackageOutputCategory? lastOutputCategory; + public OutputsPageViewModel( ISettingsManager settingsManager, IPackageFactory packageFactory, @@ -113,7 +125,7 @@ public partial class OutputsPageViewModel : PageViewModelBase // Observable predicate from SearchQuery changes var searchPredicate = this.WhenPropertyChanged(vm => vm.SearchQuery) - .Throttle(TimeSpan.FromMilliseconds(50))! + .Throttle(TimeSpan.FromMilliseconds(100))! .Select(property => searcher.GetPredicate(property.Value)) .AsObservable(); @@ -122,28 +134,37 @@ public partial class OutputsPageViewModel : PageViewModelBase .DeferUntilLoaded() .Filter(searchPredicate) .Transform(file => new OutputImageViewModel(file)) - .SortBy(vm => vm.ImageFile.CreatedAt, SortDirection.Descending) + .Sort( + SortExpressionComparer + .Descending(vm => vm.ImageFile.CreatedAt) + .ThenByDescending(vm => vm.ImageFile.FileName) + ) .Bind(Outputs) .WhenPropertyChanged(p => p.IsSelected) + .Throttle(TimeSpan.FromMilliseconds(50)) .Subscribe(_ => { NumItemsSelected = Outputs.Count(o => o.IsSelected); }); + categoriesCache.Connect().DeferUntilLoaded().Bind(Categories).Subscribe(); + settingsManager.RelayPropertyFor( this, vm => vm.ImageSize, settings => settings.OutputsImageSize, delay: TimeSpan.FromMilliseconds(250) ); - } - protected override void OnInitialLoaded() - { - RefreshCategories(); + settingsManager.RelayPropertyFor( + this, + vm => vm.ShowFolders, + settings => settings.IsOutputsTreeViewEnabled, + true + ); } - public override void OnLoaded() + protected override void OnInitialLoaded() { if (Design.IsDesignMode) return; @@ -151,14 +172,17 @@ public partial class OutputsPageViewModel : PageViewModelBase if (!settingsManager.IsLibraryDirSet) return; - base.OnLoaded(); - Directory.CreateDirectory(settingsManager.ImagesDirectory); + RefreshCategories(); + SelectedCategory ??= Categories.First(); SelectedOutputType ??= SharedOutputType.All; SearchQuery = string.Empty; ImageSize = settingsManager.Settings.OutputsImageSize; + lastOutputCategory = SelectedCategory; + + IsChangingCategory = true; var path = CanShowOutputTypes && SelectedOutputType != SharedOutputType.All @@ -169,7 +193,7 @@ public partial class OutputsPageViewModel : PageViewModelBase partial void OnSelectedCategoryChanged(PackageOutputCategory? oldValue, PackageOutputCategory? newValue) { - if (oldValue == newValue || newValue == null) + if (oldValue == newValue || oldValue == null || newValue == null) return; var path = @@ -177,11 +201,12 @@ public partial class OutputsPageViewModel : PageViewModelBase ? Path.Combine(newValue.Path, SelectedOutputType.ToString()) : SelectedCategory.Path; GetOutputs(path); + lastOutputCategory = newValue; } partial void OnSelectedOutputTypeChanged(SharedOutputType? oldValue, SharedOutputType? newValue) { - if (oldValue == newValue || newValue == null) + if (oldValue == newValue || oldValue == null || newValue == null) return; var path = @@ -262,8 +287,13 @@ public partial class OutputsPageViewModel : PageViewModelBase public void Refresh() { - Dispatcher.UIThread.Post(() => RefreshCategories()); - Dispatcher.UIThread.Post(OnLoaded); + Dispatcher.UIThread.Post(RefreshCategories); + + var path = + CanShowOutputTypes && SelectedOutputType != SharedOutputType.All + ? Path.Combine(SelectedCategory.Path, SelectedOutputType.ToString()) + : SelectedCategory.Path; + GetOutputs(path); } public async Task DeleteImage(OutputImageViewModel? item) @@ -273,8 +303,8 @@ public partial class OutputsPageViewModel : PageViewModelBase var confirmationDialog = new BetterContentDialog { - Title = "Are you sure you want to delete this image?", - Content = "This action cannot be undone.", + Title = Resources.Label_AreYouSure, + Content = Resources.Label_ActionCannotBeUndone, PrimaryButtonText = Resources.Action_Delete, SecondaryButtonText = Resources.Action_Cancel, DefaultButton = ContentDialogButton.Primary, @@ -352,8 +382,8 @@ public partial class OutputsPageViewModel : PageViewModelBase { var confirmationDialog = new BetterContentDialog { - Title = $"Are you sure you want to delete {NumItemsSelected} images?", - Content = "This action cannot be undone.", + Title = string.Format(Resources.Label_AreYouSureDeleteImages, NumItemsSelected), + Content = Resources.Label_ActionCannotBeUndone, PrimaryButtonText = Resources.Action_Delete, SecondaryButtonText = Resources.Action_Cancel, DefaultButton = ContentDialogButton.Primary, @@ -491,7 +521,7 @@ public partial class OutputsPageViewModel : PageViewModelBase } } - OnLoaded(); + Refresh(); IsConsolidating = false; } @@ -517,20 +547,58 @@ public partial class OutputsPageViewModel : PageViewModelBase return; } - var files = Directory - .EnumerateFiles(directory, "*.*", SearchOption.AllDirectories) - .Where(path => allowedExtensions.Contains(new FilePath(path).Extension)) - .Select(file => LocalImageFile.FromPath(file)) - .ToList(); - - if (files.Count == 0) + if (lastOutputCategory?.Path.Equals(directory) is not true) { OutputsCache.Clear(); + IsChangingCategory = true; } - else + + IsLoading = true; + + cancellationTokenSources.ForEach(cts => cts.Cancel()); + + Task.Run(() => { - OutputsCache.EditDiff(files); - } + var getOutputsTokenSource = new CancellationTokenSource(); + cancellationTokenSources.Add(getOutputsTokenSource); + + if (getOutputsTokenSource.IsCancellationRequested) + { + cancellationTokenSources.Remove(getOutputsTokenSource); + return; + } + + var files = Directory + .EnumerateFiles(directory, "*.*", SearchOption.AllDirectories) + .Where(file => allowedExtensions.Contains(new FilePath(file).Extension)) + .Select(file => LocalImageFile.FromPath(file)) + .ToList(); + + if (getOutputsTokenSource.IsCancellationRequested) + { + cancellationTokenSources.Remove(getOutputsTokenSource); + return; + } + + Dispatcher.UIThread.Post(() => + { + if (files.Count == 0 && OutputsCache.Count == 0) + { + IsLoading = false; + IsChangingCategory = false; + return; + } + + OutputsCache.EditDiff( + files, + (oldItem, newItem) => oldItem.AbsolutePath == newItem.AbsolutePath + ); + + IsLoading = false; + IsChangingCategory = false; + }); + cancellationTokenSources.Remove(getOutputsTokenSource); + }); } private void RefreshCategories() @@ -559,7 +627,10 @@ public partial class OutputsPageViewModel : PageViewModelBase pair.InstalledPackage.FullPath!, pair.BasePackage.OutputFolderName ), - Name = pair.InstalledPackage.DisplayName ?? "" + Name = pair.InstalledPackage.DisplayName ?? "", + SubDirectories = GetSubfolders( + Path.Combine(pair.InstalledPackage.FullPath!, pair.BasePackage.OutputFolderName) + ) } ) .ToList(); @@ -569,18 +640,37 @@ public partial class OutputsPageViewModel : PageViewModelBase new PackageOutputCategory { Path = settingsManager.ImagesDirectory, - Name = "Shared Output Folder" + Name = "Shared Output Folder", + SubDirectories = GetSubfolders(settingsManager.ImagesDirectory) } ); - packageCategories.Insert( - 1, - new PackageOutputCategory { Path = settingsManager.ImagesInferenceDirectory, Name = "Inference" } - ); + categoriesCache.EditDiff(packageCategories, (a, b) => a.Path == b.Path); + + SelectedCategory = previouslySelectedCategory ?? Categories.First(); + } + + private ObservableCollection GetSubfolders(string strPath) + { + var subfolders = new ObservableCollection(); - Categories = new ObservableCollection(packageCategories); + if (!Directory.Exists(strPath)) + return subfolders; + + var directories = Directory.EnumerateDirectories(strPath, "*", SearchOption.TopDirectoryOnly); + + foreach (var dir in directories) + { + var category = new PackageOutputCategory { Name = Path.GetFileName(dir), Path = dir }; + + if (Directory.GetDirectories(dir, "*", SearchOption.TopDirectoryOnly).Length > 0) + { + category.SubDirectories = GetSubfolders(dir); + } + + subfolders.Add(category); + } - SelectedCategory = - Categories.FirstOrDefault(x => x.Name == previouslySelectedCategory?.Name) ?? Categories.First(); + return subfolders; } } diff --git a/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs index ba815fbf..93f3534a 100644 --- a/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageCardViewModel.cs @@ -1,11 +1,16 @@ using System; using System.Collections.Generic; +using System.Collections.Immutable; +using System.Collections.Specialized; using System.Linq; +using System.Threading; using System.Threading.Tasks; +using AsyncAwaitBestPractices; using Avalonia; using Avalonia.Controls; using Avalonia.Controls.Notifications; using Avalonia.Controls.Primitives; +using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; using FluentAvalonia.UI.Controls; @@ -26,7 +31,6 @@ using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.FileInterfaces; using StabilityMatrix.Core.Models.PackageModification; using StabilityMatrix.Core.Models.Packages; -using StabilityMatrix.Core.Models.Settings; using StabilityMatrix.Core.Processes; using StabilityMatrix.Core.Services; @@ -34,14 +38,17 @@ namespace StabilityMatrix.Avalonia.ViewModels.PackageManager; [ManagedService] [Transient] -public partial class PackageCardViewModel : ProgressViewModel +public partial class PackageCardViewModel( + ILogger logger, + IPackageFactory packageFactory, + INotificationService notificationService, + ISettingsManager settingsManager, + INavigationService navigationService, + ServiceManager vmFactory, + RunningPackageService runningPackageService +) : ProgressViewModel { - private readonly ILogger logger; - private readonly IPackageFactory packageFactory; - private readonly INotificationService notificationService; - private readonly ISettingsManager settingsManager; - private readonly INavigationService navigationService; - private readonly ServiceManager vmFactory; + private string webUiUrl = string.Empty; [ObservableProperty] private InstalledPackage? package; @@ -82,21 +89,31 @@ public partial class PackageCardViewModel : ProgressViewModel [ObservableProperty] private bool canUseExtensions; - public PackageCardViewModel( - ILogger logger, - IPackageFactory packageFactory, - INotificationService notificationService, - ISettingsManager settingsManager, - INavigationService navigationService, - ServiceManager vmFactory - ) + [ObservableProperty] + private bool isRunning; + + [ObservableProperty] + private bool showWebUiButton; + + private void RunningPackagesOnCollectionChanged(object? sender, NotifyCollectionChangedEventArgs e) { - this.logger = logger; - this.packageFactory = packageFactory; - this.notificationService = notificationService; - this.settingsManager = settingsManager; - this.navigationService = navigationService; - this.vmFactory = vmFactory; + if (runningPackageService.RunningPackages.Select(x => x.Value) is not { } runningPackages) + return; + + var runningViewModel = runningPackages.FirstOrDefault( + x => x.RunningPackage.InstalledPackage.Id == Package?.Id + ); + if (runningViewModel is not null) + { + IsRunning = true; + runningViewModel.RunningPackage.BasePackage.Exited += BasePackageOnExited; + runningViewModel.RunningPackage.BasePackage.StartupComplete += RunningPackageOnStartupComplete; + } + else if (runningViewModel is null && IsRunning) + { + IsRunning = false; + ShowWebUiButton = false; + } } partial void OnPackageChanged(InstalledPackage? value) @@ -127,11 +144,30 @@ public partial class PackageCardViewModel : ProgressViewModel UseSharedOutput = Package?.UseSharedOutputFolder ?? false; CanUseSharedOutput = basePackage?.SharedOutputFolders != null; CanUseExtensions = basePackage?.SupportsExtensions ?? false; + + runningPackageService.RunningPackages.CollectionChanged += RunningPackagesOnCollectionChanged; + EventManager.Instance.PackageRelaunchRequested += InstanceOnPackageRelaunchRequested; } } + private void InstanceOnPackageRelaunchRequested(object? sender, InstalledPackage e) + { + if (e.Id != Package?.Id) + return; + + navigationService.GoBack(); + Launch().SafeFireAndForget(); + } + public override async Task OnLoadedAsync() { + if (Design.IsDesignMode && Package?.DisplayName == "Running Comfy") + { + IsRunning = true; + IsUpdateAvailable = true; + ShowWebUiButton = true; + } + if (Design.IsDesignMode || !settingsManager.IsLibraryDirSet || Package is not { } currentPackage) return; @@ -160,18 +196,122 @@ public partial class PackageCardViewModel : ProgressViewModel } IsUpdateAvailable = await HasUpdate(); + + if ( + Package != null + && !IsRunning + && runningPackageService.RunningPackages.TryGetValue(Package.Id, out var runningPackageVm) + ) + { + IsRunning = true; + runningPackageVm.RunningPackage.BasePackage.Exited += BasePackageOnExited; + runningPackageVm.RunningPackage.BasePackage.StartupComplete += + RunningPackageOnStartupComplete; + webUiUrl = runningPackageVm.WebUiUrl; + ShowWebUiButton = !string.IsNullOrWhiteSpace(webUiUrl); + } } } - public void Launch() + public override void OnUnloaded() + { + EventManager.Instance.PackageRelaunchRequested -= InstanceOnPackageRelaunchRequested; + runningPackageService.RunningPackages.CollectionChanged -= RunningPackagesOnCollectionChanged; + } + + public async Task Launch() { if (Package == null) return; - settingsManager.Transaction(s => s.ActiveInstalledPackageId = Package.Id); + var packagePair = await runningPackageService.StartPackage(Package); + + if (packagePair != null) + { + IsRunning = true; + + packagePair.BasePackage.Exited += BasePackageOnExited; + packagePair.BasePackage.StartupComplete += RunningPackageOnStartupComplete; + + var vm = runningPackageService.GetRunningPackageViewModel(packagePair.InstalledPackage.Id); + if (vm != null) + { + navigationService.NavigateTo(vm, new BetterEntranceNavigationTransition()); + } + } + + // settingsManager.Transaction(s => s.ActiveInstalledPackageId = Package.Id); + // + // navigationService.NavigateTo(new BetterDrillInNavigationTransition()); + // EventManager.Instance.OnPackageLaunchRequested(Package.Id); + } + + public void NavToConsole() + { + if (Package == null) + return; + + var vm = runningPackageService.GetRunningPackageViewModel(Package.Id); + if (vm != null) + { + navigationService.NavigateTo(vm, new BetterEntranceNavigationTransition()); + } + } + + public void LaunchWebUi() + { + if (string.IsNullOrEmpty(webUiUrl)) + return; + + notificationService.TryAsync( + Task.Run(() => ProcessRunner.OpenUrl(webUiUrl)), + "Failed to open URL", + $"{webUiUrl}" + ); + } + + private void BasePackageOnExited(object? sender, int exitCode) + { + Dispatcher + .UIThread.InvokeAsync(async () => + { + logger.LogTrace("Process exited ({Code}) at {Time:g}", exitCode, DateTimeOffset.Now); + + // Need to wait for streams to finish before detaching handlers + if (sender is BaseGitPackage { VenvRunner: not null } package) + { + var process = package.VenvRunner.Process; + if (process is not null) + { + // Max 5 seconds + var ct = new CancellationTokenSource(5000).Token; + try + { + await process.WaitUntilOutputEOF(ct); + } + catch (OperationCanceledException e) + { + logger.LogWarning("Waiting for process EOF timed out: {Message}", e.Message); + } + } + } + + // Detach handlers + if (sender is BasePackage basePackage) + { + basePackage.Exited -= BasePackageOnExited; + basePackage.StartupComplete -= RunningPackageOnStartupComplete; + } + + if (Package?.Id != null) + { + runningPackageService.RunningPackages.Remove(Package.Id); + } - navigationService.NavigateTo(new BetterDrillInNavigationTransition()); - EventManager.Instance.OnPackageLaunchRequested(Package.Id); + IsRunning = false; + ShowWebUiButton = false; + }) + .SafeFireAndForget(); } public async Task Uninstall() @@ -181,13 +321,17 @@ public partial class PackageCardViewModel : ProgressViewModel return; } - var dialog = new ContentDialog + var dialogViewModel = vmFactory.Get(vm => { - Title = Resources.Label_ConfirmDelete, - Content = Resources.Text_PackageUninstall_Details, - PrimaryButtonText = Resources.Action_OK, - CloseButtonText = Resources.Action_Cancel, - DefaultButton = ContentDialogButton.Primary + vm.ExpectedPackageName = Package?.DisplayName; + }); + + var dialog = new BetterContentDialog + { + Content = dialogViewModel, + IsPrimaryButtonEnabled = false, + IsSecondaryButtonEnabled = false, + IsFooterVisible = false, }; var result = await dialog.ShowAsync(); @@ -282,6 +426,21 @@ public partial class PackageCardViewModel : ProgressViewModel versionOptions.CommitHash = latest.Sha; } + var confirmationDialog = new BetterContentDialog + { + Title = Resources.Label_AreYouSure, + Content = + $"{Package.DisplayName} will be updated to the latest version ({versionOptions.GetReadableVersionString()})", + PrimaryButtonText = Resources.Action_Continue, + SecondaryButtonText = Resources.Action_Cancel, + DefaultButton = ContentDialogButton.Primary, + IsSecondaryButtonEnabled = true, + }; + + var dialogResult = await confirmationDialog.ShowAsync(); + if (dialogResult != ContentDialogResult.Primary) + return; + var updatePackageStep = new UpdatePackageStep( settingsManager, Package, @@ -435,6 +594,64 @@ public partial class PackageCardViewModel : ProgressViewModel ProcessRunner.OpenUrl(basePackage.GithubUrl); } + [RelayCommand] + private async Task Stop() + { + if (Package is null) + return; + + await runningPackageService.StopPackage(Package.Id); + IsRunning = false; + ShowWebUiButton = false; + } + + [RelayCommand] + private async Task Restart() + { + await Stop(); + await Launch(); + } + + [RelayCommand] + private async Task ShowLaunchOptions() + { + var basePackage = packageFactory.FindPackageByName(Package?.PackageName); + if (basePackage == null) + { + logger.LogWarning("Package {Name} not found", Package?.PackageName); + return; + } + + var viewModel = vmFactory.Get(); + viewModel.Cards = LaunchOptionCard + .FromDefinitions(basePackage.LaunchOptions, Package?.LaunchArgs ?? []) + .ToImmutableArray(); + + logger.LogDebug("Launching config dialog with cards: {CardsCount}", viewModel.Cards.Count); + + var dialog = new BetterContentDialog + { + ContentVerticalScrollBarVisibility = ScrollBarVisibility.Disabled, + IsPrimaryButtonEnabled = true, + PrimaryButtonText = Resources.Action_Save, + CloseButtonText = Resources.Action_Cancel, + FullSizeDesired = true, + DefaultButton = ContentDialogButton.Primary, + ContentMargin = new Thickness(32, 16), + Padding = new Thickness(0, 16), + Content = new LaunchOptionsDialog { DataContext = viewModel, } + }; + + var result = await dialog.ShowAsync(); + + if (result == ContentDialogResult.Primary && Package != null) + { + // Save config + var args = viewModel.AsLaunchArgs(); + settingsManager.SaveLaunchArgs(Package.Id, args); + } + } + private async Task HasUpdate() { if (Package == null || IsUnknownPackage || Design.IsDesignMode) @@ -565,4 +782,10 @@ public partial class PackageCardViewModel : ProgressViewModel IsSharedModelConfig = false; } } + + private void RunningPackageOnStartupComplete(object? sender, string e) + { + webUiUrl = e.Replace("0.0.0.0", "127.0.0.1"); + ShowWebUiButton = !string.IsNullOrWhiteSpace(webUiUrl); + } } diff --git a/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageInstallDetailViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageInstallDetailViewModel.cs index 71e14b31..9172d5c5 100644 --- a/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageInstallDetailViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/PackageManager/PackageInstallDetailViewModel.cs @@ -281,7 +281,7 @@ public partial class PackageInstallDetailViewModel( SelectedVersion = !IsReleaseMode ? AvailableVersions?.FirstOrDefault(x => x.TagName == SelectedPackage.MainBranch) ?? AvailableVersions?.FirstOrDefault() - : AvailableVersions?.FirstOrDefault(); + : AvailableVersions?.FirstOrDefault(v => !v.IsPrerelease); CanInstall = !ShowDuplicateWarning; } diff --git a/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs index 6066fd02..9d59df05 100644 --- a/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/PackageManagerViewModel.cs @@ -6,8 +6,6 @@ using System.Linq; using System.Threading.Tasks; using AsyncAwaitBestPractices; using Avalonia.Controls; -using Avalonia.Controls.Notifications; -using Avalonia.Controls.Primitives; using Avalonia.Threading; using CommunityToolkit.Mvvm.Input; using DynamicData; @@ -15,18 +13,15 @@ using DynamicData.Binding; using FluentAvalonia.UI.Controls; using Microsoft.Extensions.Logging; using StabilityMatrix.Avalonia.Animations; -using StabilityMatrix.Avalonia.Controls; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Avalonia.ViewModels.Base; -using StabilityMatrix.Avalonia.ViewModels.Dialogs; using StabilityMatrix.Avalonia.ViewModels.PackageManager; using StabilityMatrix.Avalonia.Views; -using StabilityMatrix.Avalonia.Views.Dialogs; using StabilityMatrix.Core.Attributes; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.FileInterfaces; -using StabilityMatrix.Core.Models.PackageModification; using StabilityMatrix.Core.Models.Packages; using StabilityMatrix.Core.Services; using Symbol = FluentIcons.Common.Symbol; @@ -48,7 +43,7 @@ public partial class PackageManagerViewModel : PageViewModelBase private readonly INavigationService packageNavigationService; private readonly ILogger logger; - public override string Title => "Packages"; + public override string Title => Resources.Label_Packages; public override IconSource IconSource => new SymbolIconSource { Symbol = Symbol.Box, IsFilled = true }; /// @@ -84,6 +79,7 @@ public partial class PackageManagerViewModel : PageViewModelBase this.logger = logger; EventManager.Instance.InstalledPackagesChanged += OnInstalledPackagesChanged; + EventManager.Instance.OneClickInstallFinished += OnOneClickInstallFinished; var installed = installedPackages.Connect(); var unknown = unknownInstalledPackages.Connect(); @@ -107,6 +103,11 @@ public partial class PackageManagerViewModel : PageViewModelBase timer.Tick += async (_, _) => await CheckPackagesForUpdates(); } + private void OnOneClickInstallFinished(object? sender, bool e) + { + OnLoadedAsync().SafeFireAndForget(); + } + public void SetPackages(IEnumerable packages) { installedPackages.Edit(s => s.Load(packages)); @@ -119,7 +120,7 @@ public partial class PackageManagerViewModel : PageViewModelBase public override async Task OnLoadedAsync() { - if (Design.IsDesignMode) + if (Design.IsDesignMode || !settingsManager.IsLibraryDirSet) return; installedPackages.EditDiff(settingsManager.Settings.InstalledPackages, InstalledPackage.Comparer); diff --git a/StabilityMatrix.Avalonia/ViewModels/RunningPackageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/RunningPackageViewModel.cs new file mode 100644 index 00000000..38011e04 --- /dev/null +++ b/StabilityMatrix.Avalonia/ViewModels/RunningPackageViewModel.cs @@ -0,0 +1,161 @@ +using System; +using System.Threading.Tasks; +using CommunityToolkit.Mvvm.ComponentModel; +using CommunityToolkit.Mvvm.Input; +using FluentAvalonia.UI.Controls; +using StabilityMatrix.Avalonia.Services; +using StabilityMatrix.Avalonia.ViewModels.Base; +using StabilityMatrix.Avalonia.Views; +using StabilityMatrix.Core.Attributes; +using StabilityMatrix.Core.Helper; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Packages; +using StabilityMatrix.Core.Processes; +using StabilityMatrix.Core.Services; +using SymbolIconSource = FluentIcons.Avalonia.Fluent.SymbolIconSource; +using TeachingTip = StabilityMatrix.Core.Models.Settings.TeachingTip; + +namespace StabilityMatrix.Avalonia.ViewModels; + +[View(typeof(ConsoleOutputPage))] +public partial class RunningPackageViewModel : PageViewModelBase, IDisposable, IAsyncDisposable +{ + private readonly ISettingsManager settingsManager; + private readonly INotificationService notificationService; + private readonly RunningPackageService runningPackageService; + + public PackagePair RunningPackage { get; } + public ConsoleViewModel Console { get; } + public override string Title => RunningPackage.InstalledPackage.PackageName ?? "Running Package"; + public override IconSource IconSource => new SymbolIconSource(); + + [ObservableProperty] + private bool autoScrollToEnd; + + [ObservableProperty] + private bool showWebUiButton; + + [ObservableProperty] + private string webUiUrl = string.Empty; + + [ObservableProperty] + private bool isRunning = true; + + [ObservableProperty] + private string consoleInput = string.Empty; + + [ObservableProperty] + private bool showWebUiTeachingTip; + + /// + public RunningPackageViewModel( + ISettingsManager settingsManager, + INotificationService notificationService, + RunningPackageService runningPackageService, + PackagePair runningPackage, + ConsoleViewModel console + ) + { + this.settingsManager = settingsManager; + this.notificationService = notificationService; + this.runningPackageService = runningPackageService; + + RunningPackage = runningPackage; + Console = console; + Console.Document.LineCountChanged += DocumentOnLineCountChanged; + RunningPackage.BasePackage.StartupComplete += BasePackageOnStartupComplete; + RunningPackage.BasePackage.Exited += BasePackageOnExited; + + settingsManager.RelayPropertyFor( + this, + vm => vm.AutoScrollToEnd, + settings => settings.AutoScrollLaunchConsoleToEnd, + true + ); + } + + private void BasePackageOnExited(object? sender, int e) + { + IsRunning = false; + ShowWebUiButton = false; + Console.Document.LineCountChanged -= DocumentOnLineCountChanged; + RunningPackage.BasePackage.StartupComplete -= BasePackageOnStartupComplete; + RunningPackage.BasePackage.Exited -= BasePackageOnExited; + runningPackageService.RunningPackages.Remove(RunningPackage.InstalledPackage.Id); + } + + private void BasePackageOnStartupComplete(object? sender, string url) + { + WebUiUrl = url.Replace("0.0.0.0", "127.0.0.1"); + ShowWebUiButton = !string.IsNullOrWhiteSpace(WebUiUrl); + + if (settingsManager.Settings.SeenTeachingTips.Contains(TeachingTip.WebUiButtonMovedTip)) + return; + + ShowWebUiTeachingTip = true; + settingsManager.Transaction(s => s.SeenTeachingTips.Add(TeachingTip.WebUiButtonMovedTip)); + } + + private void DocumentOnLineCountChanged(object? sender, EventArgs e) + { + if (AutoScrollToEnd) + { + EventManager.Instance.OnScrollToBottomRequested(); + } + } + + [RelayCommand] + private void LaunchPackage() + { + EventManager.Instance.OnPackageRelaunchRequested(RunningPackage.InstalledPackage); + } + + [RelayCommand] + private async Task Stop() + { + await runningPackageService.StopPackage(RunningPackage.InstalledPackage.Id); + Console.PostLine($"{Environment.NewLine}Stopped process at {DateTimeOffset.Now}"); + await Console.StopUpdatesAsync(); + IsRunning = false; + } + + [RelayCommand] + private void LaunchWebUi() + { + if (string.IsNullOrEmpty(WebUiUrl)) + return; + + notificationService.TryAsync( + Task.Run(() => ProcessRunner.OpenUrl(WebUiUrl)), + "Failed to open URL", + $"{WebUiUrl}" + ); + } + + [RelayCommand] + private async Task SendToConsole() + { + Console.PostLine(ConsoleInput); + if (RunningPackage?.BasePackage is BaseGitPackage gitPackage) + { + var venv = gitPackage.VenvRunner; + var process = venv?.Process; + if (process is not null) + { + await process.StandardInput.WriteLineAsync(ConsoleInput); + } + } + + ConsoleInput = string.Empty; + } + + public void Dispose() + { + Console.Dispose(); + } + + public async ValueTask DisposeAsync() + { + await Console.DisposeAsync(); + } +} diff --git a/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs index cf46f2c6..0da8ab4c 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs @@ -14,6 +14,7 @@ using System.Text; using System.Text.Json; using System.Threading.Tasks; using AsyncAwaitBestPractices; +using AsyncImageLoader; using Avalonia; using Avalonia.Controls.Notifications; using Avalonia.Controls.Primitives; @@ -784,7 +785,10 @@ public partial class MainSettingsViewModel : PageViewModelBase [ new CommandItem(DebugFindLocalModelFromIndexCommand), new CommandItem(DebugExtractDmgCommand), - new CommandItem(DebugShowNativeNotificationCommand) + new CommandItem(DebugShowNativeNotificationCommand), + new CommandItem(DebugClearImageCacheCommand), + new CommandItem(DebugGCCollectCommand), + new CommandItem(DebugExtractImagePromptsToTxtCommand) ]; [RelayCommand] @@ -899,6 +903,60 @@ public partial class MainSettingsViewModel : PageViewModelBase ); } + [RelayCommand] + private void DebugClearImageCache() + { + if (ImageLoader.AsyncImageLoader is FallbackRamCachedWebImageLoader loader) + { + loader.ClearCache(); + } + } + + [RelayCommand] + private void DebugGCCollect() + { + GC.Collect(); + } + + [RelayCommand] + private async Task DebugExtractImagePromptsToTxt() + { + // Choose images + var provider = App.StorageProvider; + var files = await provider.OpenFilePickerAsync(new FilePickerOpenOptions { AllowMultiple = true }); + + if (files.Count == 0) + return; + + var images = await Task.Run( + () => files.Select(f => LocalImageFile.FromPath(f.TryGetLocalPath()!)).ToList() + ); + + var successfulFiles = new List(); + + foreach (var localImage in images) + { + var imageFile = new FilePath(localImage.AbsolutePath); + + // Write a txt with the same name as the image + var txtFile = imageFile.WithName(imageFile.NameWithoutExtension + ".txt"); + + // Read metadata + if (localImage.GenerationParameters?.PositivePrompt is { } positivePrompt) + { + await File.WriteAllTextAsync(txtFile, positivePrompt); + + successfulFiles.Add(localImage); + } + } + + notificationService.Show( + "Extracted prompts", + $"Extracted prompts from {successfulFiles.Count}/{images.Count} images.", + NotificationType.Success + ); + } + #endregion #region Info Section diff --git a/StabilityMatrix.Avalonia/Views/CivitAiBrowserPage.axaml b/StabilityMatrix.Avalonia/Views/CivitAiBrowserPage.axaml index ff25d806..25755a05 100644 --- a/StabilityMatrix.Avalonia/Views/CivitAiBrowserPage.axaml +++ b/StabilityMatrix.Avalonia/Views/CivitAiBrowserPage.axaml @@ -15,6 +15,8 @@ xmlns:converters="clr-namespace:StabilityMatrix.Avalonia.Converters" xmlns:asyncImageLoader="clr-namespace:AsyncImageLoader;assembly=AsyncImageLoader.Avalonia" xmlns:helpers="clr-namespace:StabilityMatrix.Avalonia.Helpers" + xmlns:controls1="clr-namespace:Avalonia.Labs.Controls;assembly=Avalonia.Labs.Controls" + xmlns:system="clr-namespace:System;assembly=System.Runtime" d:DataContext="{x:Static designData:DesignData.CivitAiBrowserViewModel}" d:DesignHeight="700" d:DesignWidth="800" @@ -70,6 +72,7 @@ + False - - - + + Stretch="UniformToFill"/> - + Value="{Binding CivitModel.ModelVersionStats.Rating}" /> @@ -336,7 +338,7 @@ + Text="{Binding CivitModel.Stats.ThumbsUpCount, Converter={StaticResource KiloFormatterConverter}}" /> - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + - - - + - + + + - - - + IsEnabled="{Binding !IsLoading}" + Label="{x:Static lang:Resources.Action_Refresh}"> + + + + + + @@ -160,11 +149,22 @@ + + - + @@ -206,7 +206,8 @@ IconSource="Delete" Text="{x:Static lang:Resources.Action_Delete}" /> - + - - - + + + - - - + + + - - - - - - - + + + + + + + - - + + + - - - - - + + + - - - - - - - - - - - + + + - - - - - - - - - + + + + + + + + + + diff --git a/StabilityMatrix.Avalonia/Views/PackageManagerPage.axaml b/StabilityMatrix.Avalonia/Views/PackageManagerPage.axaml index 456c3697..81c912e6 100644 --- a/StabilityMatrix.Avalonia/Views/PackageManagerPage.axaml +++ b/StabilityMatrix.Avalonia/Views/PackageManagerPage.axaml @@ -10,15 +10,23 @@ xmlns:icons="clr-namespace:Projektanker.Icons.Avalonia;assembly=Projektanker.Icons.Avalonia" xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages" xmlns:system="clr-namespace:System;assembly=System.Runtime" - mc:Ignorable="d" d:DesignWidth="800" d:DesignHeight="450" + xmlns:converters="clr-namespace:StabilityMatrix.Avalonia.Converters" + xmlns:avalonia="clr-namespace:SpacedGridControl.Avalonia;assembly=SpacedGridControl.Avalonia" + xmlns:markupExtensions="clr-namespace:StabilityMatrix.Avalonia.MarkupExtensions" + mc:Ignorable="d" d:DesignWidth="1000" d:DesignHeight="450" x:DataType="viewModels:PackageManagerViewModel" x:CompileBindings="True" d:DataContext="{x:Static designData:DesignData.PackageManagerViewModel}" x:Class="StabilityMatrix.Avalonia.Views.PackageManagerPage"> + + + + @@ -26,185 +34,34 @@ - - - - - - - - + + + + + + + + - - @@ -215,99 +72,340 @@ TextWrapping="Wrap" Text="{x:Static lang:Resources.Label_UnknownPackage}" /> + + + + + - - + + + + + + + + + + - - + - - + + + + + + + - - + + + + + + + + + - + + IsVisible="{Binding IsProgressVisible}" + RowDefinitions="Auto, *"> + + public record ModelConnections(string Name) { + public ModelConnections(ModelConnections other) + { + Name = other.Name; + Model = other.Model; + VAE = other.VAE; + Clip = other.Clip; + Conditioning = other.Conditioning; + } + public ModelNodeConnection? Model { get; set; } public VAENodeConnection? VAE { get; set; } diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs index 286dff0b..c9a378b6 100644 --- a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs +++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs @@ -7,6 +7,7 @@ using StabilityMatrix.Core.Attributes; using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; using StabilityMatrix.Core.Models.Database; +using StabilityMatrix.Core.Models.Inference; namespace StabilityMatrix.Core.Models.Api.Comfy.Nodes; @@ -131,23 +132,35 @@ public class ComfyNodeBuilder public int StopAtClipLayer { get; init; } = -1; } - public static NamedComfyNode LatentFromBatch( - string name, - LatentNodeConnection samples, - int batchIndex, - int length - ) + public record LatentFromBatch : ComfyTypedNodeBase { - return new NamedComfyNode(name) - { - ClassType = "LatentFromBatch", - Inputs = new Dictionary - { - ["samples"] = samples.Data, - ["batch_index"] = batchIndex, - ["length"] = length, - } - }; + public required LatentNodeConnection Samples { get; init; } + + [Range(0, 63)] + public int BatchIndex { get; init; } = 0; + + [Range(1, 64)] + public int Length { get; init; } = 1; + } + + public record LatentBlend : ComfyTypedNodeBase + { + public required LatentNodeConnection Samples1 { get; init; } + + public required LatentNodeConnection Samples2 { get; init; } + + [Range(0d, 1d)] + public double BlendFactor { get; init; } = 0.5; + } + + public record ModelMergeSimple : ComfyTypedNodeBase + { + public required ModelNodeConnection Model1 { get; init; } + + public required ModelNodeConnection Model2 { get; init; } + + [Range(0d, 1d)] + public double Ratio { get; init; } = 1; } public static NamedComfyNode ImageUpscaleWithModel( @@ -332,7 +345,7 @@ public class ComfyNodeBuilder [TypedNodeOptions( Name = "Inference_Core_PromptExpansion", - RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes"] + RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.2.0"] )] public record PromptExpansion : ComfyTypedNodeBase { @@ -342,6 +355,34 @@ public class ComfyNodeBuilder public bool LogPrompt { get; init; } } + [TypedNodeOptions( + Name = "Inference_Core_AIO_Preprocessor", + RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.2.0"] + )] + public record AIOPreprocessor : ComfyTypedNodeBase + { + public required ImageNodeConnection Image { get; init; } + + public required string Preprocessor { get; init; } + + [Range(64, 2048)] + public int Resolution { get; init; } = 512; + } + + [TypedNodeOptions( + Name = "Inference_Core_ReferenceOnlySimple", + RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.3.0"] + )] + public record ReferenceOnlySimple : ComfyTypedNodeBase + { + public required ModelNodeConnection Model { get; init; } + + public required LatentNodeConnection Reference { get; init; } + + [Range(1, 64)] + public int BatchSize { get; init; } = 1; + } + public ImageNodeConnection Lambda_LatentToImage(LatentNodeConnection latent, VAENodeConnection vae) { var name = GetUniqueName("VAEDecode"); @@ -818,6 +859,14 @@ public class ComfyNodeBuilder public ModelConnections Base => Models["Base"]; public ModelConnections Refiner => Models["Refiner"]; + public Dictionary SamplerTemporaryArgs { get; } = new(); + + public ModuleApplyStepTemporaryArgs? BaseSamplerTemporaryArgs + { + get => SamplerTemporaryArgs.GetValueOrDefault("Base"); + set => SamplerTemporaryArgs["Base"] = value; + } + public PrimaryNodeConnection? Primary { get; set; } public VAENodeConnection? PrimaryVAE { get; set; } public Size PrimarySize { get; set; } diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/NodeDictionary.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/NodeDictionary.cs index 5fbe4464..da088767 100644 --- a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/NodeDictionary.cs +++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/NodeDictionary.cs @@ -1,10 +1,12 @@ using System.ComponentModel; using System.Reflection; using System.Text.Json.Serialization; +using KGySoft.CoreLibraries; using OneOf; using StabilityMatrix.Core.Attributes; using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; +using StabilityMatrix.Core.Models.Packages.Extensions; namespace StabilityMatrix.Core.Models.Api.Comfy.Nodes; @@ -19,7 +21,10 @@ public class NodeDictionary : Dictionary /// When inserting TypedNodes, this holds a mapping of ClassType to required extensions /// [JsonIgnore] - public Dictionary ClassTypeRequiredExtensions { get; } = new(); + public Dictionary ClassTypeRequiredExtensions { get; } = new(); + + public IEnumerable RequiredExtensions => + ClassTypeRequiredExtensions.Values.SelectMany(x => x); /// /// Finds a unique node name given a base name, by appending _2, _3, etc. @@ -63,7 +68,11 @@ public class NodeDictionary : Dictionary { if (options.RequiredExtensions != null) { - ClassTypeRequiredExtensions[namedNode.ClassType] = options.RequiredExtensions; + ClassTypeRequiredExtensions.AddOrUpdate( + namedNode.ClassType, + _ => options.GetRequiredExtensionSpecifiers().ToArray(), + (_, specifiers) => options.GetRequiredExtensionSpecifiers().Concat(specifiers).ToArray() + ); } } diff --git a/StabilityMatrix.Core/Models/DownloadPackageVersionOptions.cs b/StabilityMatrix.Core/Models/DownloadPackageVersionOptions.cs index c8614ada..da41388a 100644 --- a/StabilityMatrix.Core/Models/DownloadPackageVersionOptions.cs +++ b/StabilityMatrix.Core/Models/DownloadPackageVersionOptions.cs @@ -7,4 +7,7 @@ public class DownloadPackageVersionOptions public string? VersionTag { get; set; } public bool IsLatest { get; set; } public bool IsPrerelease { get; set; } + + public string GetReadableVersionString() => + !string.IsNullOrWhiteSpace(VersionTag) ? VersionTag : $"{BranchName}@{CommitHash?[..7]}"; } diff --git a/StabilityMatrix.Core/Models/HybridModelFile.cs b/StabilityMatrix.Core/Models/HybridModelFile.cs index 5e0d7275..8e21cb39 100644 --- a/StabilityMatrix.Core/Models/HybridModelFile.cs +++ b/StabilityMatrix.Core/Models/HybridModelFile.cs @@ -1,5 +1,6 @@ using System.Diagnostics.CodeAnalysis; using System.Text.Json.Serialization; +using StabilityMatrix.Core.Helper; using StabilityMatrix.Core.Models.Database; namespace StabilityMatrix.Core.Models; @@ -20,7 +21,7 @@ public record HybridModelFile /// public static HybridModelFile None { get; } = FromRemote("@none"); - private string? RemoteName { get; init; } + public string? RemoteName { get; init; } public LocalModelFile? Local { get; init; } @@ -67,6 +68,11 @@ public record HybridModelFile return "Default"; } + if (ReferenceEquals(this, RemoteModels.ControlNetReferenceOnlyModel)) + { + return "Reference Only"; + } + var fileName = Path.GetFileNameWithoutExtension(RelativePath); if ( diff --git a/StabilityMatrix.Core/Models/Inference/ModuleApplyStepTemporaryArgs.cs b/StabilityMatrix.Core/Models/Inference/ModuleApplyStepTemporaryArgs.cs new file mode 100644 index 00000000..e76ea875 --- /dev/null +++ b/StabilityMatrix.Core/Models/Inference/ModuleApplyStepTemporaryArgs.cs @@ -0,0 +1,46 @@ +using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; + +namespace StabilityMatrix.Core.Models.Inference; + +public class ModuleApplyStepTemporaryArgs +{ + /// + /// Temporary Primary apply step, used by ControlNet ReferenceOnly which changes the latent. + /// + public PrimaryNodeConnection? Primary { get; set; } + + public VAENodeConnection? PrimaryVAE { get; set; } + + /// + /// Used by Reference-Only ControlNet to indicate that has been batched. + /// + public bool IsPrimaryTempBatched { get; set; } + + /// + /// When is true, this is the index of the temp batch to pick after sampling. + /// + public int PrimaryTempBatchPickIndex { get; set; } + + public Dictionary Models { get; set; } = + new() { ["Base"] = new ModelConnections("Base"), ["Refiner"] = new ModelConnections("Refiner") }; + + public ModelConnections Base => Models["Base"]; + public ModelConnections Refiner => Models["Refiner"]; + + public ConditioningConnections GetRefinerOrBaseConditioning() + { + return Refiner.Conditioning + ?? Base.Conditioning + ?? throw new NullReferenceException("No Refiner or Base Conditioning"); + } + + public ModelNodeConnection GetRefinerOrBaseModel() + { + return Refiner.Model ?? Base.Model ?? throw new NullReferenceException("No Refiner or Base Model"); + } + + public VAENodeConnection GetDefaultVAE() + { + return PrimaryVAE ?? Refiner.VAE ?? Base.VAE ?? throw new NullReferenceException("No VAE"); + } +} diff --git a/StabilityMatrix.Core/Models/InferenceRunCustomPromptEventArgs.cs b/StabilityMatrix.Core/Models/InferenceRunCustomPromptEventArgs.cs new file mode 100644 index 00000000..2c09c8de --- /dev/null +++ b/StabilityMatrix.Core/Models/InferenceRunCustomPromptEventArgs.cs @@ -0,0 +1,14 @@ +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; + +namespace StabilityMatrix.Core.Models; + +public class InferenceQueueCustomPromptEventArgs : EventArgs +{ + public ComfyNodeBuilder Builder { get; } = new(); + + public NodeDictionary Nodes => Builder.Nodes; + + public long? SeedOverride { get; init; } + + public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = []; +} diff --git a/StabilityMatrix.Core/Models/PackageDifficulty.cs b/StabilityMatrix.Core/Models/PackageDifficulty.cs index 51eacc46..03cc6232 100644 --- a/StabilityMatrix.Core/Models/PackageDifficulty.cs +++ b/StabilityMatrix.Core/Models/PackageDifficulty.cs @@ -2,11 +2,13 @@ public enum PackageDifficulty { + ReallyRecommended = -1, Recommended = 0, - Simple = 1, - Advanced = 2, - Expert = 3, - Nightmare = 4, - UltraNightmare = 5, + InferenceCompatible = 1, + Simple = 2, + Advanced = 3, + Expert = 4, + Nightmare = 5, + UltraNightmare = 6, Impossible = 999 } diff --git a/StabilityMatrix.Core/Models/PackageModification/IPackageModificationRunner.cs b/StabilityMatrix.Core/Models/PackageModification/IPackageModificationRunner.cs index 0c99500b..4a4a0e0b 100644 --- a/StabilityMatrix.Core/Models/PackageModification/IPackageModificationRunner.cs +++ b/StabilityMatrix.Core/Models/PackageModification/IPackageModificationRunner.cs @@ -32,7 +32,7 @@ public interface IPackageModificationRunner string? ModificationCompleteTitle { get; init; } - string? ModificationCompleteMessage { get; init; } + string ModificationCompleteMessage { get; init; } string? ModificationFailedTitle { get; init; } diff --git a/StabilityMatrix.Core/Models/PackageModification/PackageModificationRunner.cs b/StabilityMatrix.Core/Models/PackageModification/PackageModificationRunner.cs index 902717c5..6d949f92 100644 --- a/StabilityMatrix.Core/Models/PackageModification/PackageModificationRunner.cs +++ b/StabilityMatrix.Core/Models/PackageModification/PackageModificationRunner.cs @@ -80,7 +80,7 @@ public class PackageModificationRunner : IPackageModificationRunner public string? ModificationCompleteTitle { get; init; } = "Install Complete"; - public string? ModificationCompleteMessage { get; init; } + public required string ModificationCompleteMessage { get; init; } public string? ModificationFailedTitle { get; init; } = "Install Failed"; diff --git a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs index df78aa1f..f823efeb 100644 --- a/StabilityMatrix.Core/Models/Packages/A3WebUI.cs +++ b/StabilityMatrix.Core/Models/Packages/A3WebUI.cs @@ -97,6 +97,13 @@ public class A3WebUI( DefaultValue = "7860", Options = ["--port"] }, + new LaunchOptionDefinition + { + Name = "Share", + Type = LaunchOptionType.Bool, + Description = "Set whether to share on Gradio", + Options = { "--share" } + }, new() { Name = "VRAM", diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs index 45358a5e..da47bbb6 100644 --- a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs +++ b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs @@ -1,6 +1,5 @@ using System.Diagnostics; using System.Text.Json; -using System.Text.Json.Serialization; using System.Text.RegularExpressions; using NLog; using StabilityMatrix.Core.Attributes; @@ -44,7 +43,7 @@ public class ComfyUI( public override bool ShouldIgnoreReleases => true; public override bool IsInferenceCompatible => true; public override string OutputFolderName => "output"; - public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Advanced; + public override PackageDifficulty InstallerSortOrder => PackageDifficulty.InferenceCompatible; public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.Configuration; @@ -221,7 +220,7 @@ public class ComfyUI( { TorchVersion.Cpu => "cpu", TorchVersion.Cuda => "cu121", - TorchVersion.Rocm => "rocm5.6", + TorchVersion.Rocm => "rocm5.7", _ => throw new ArgumentOutOfRangeException( nameof(torchVersion), @@ -232,9 +231,14 @@ public class ComfyUI( ) }; - if (torchVersion == TorchVersion.Cuda) + switch (torchVersion) { - pipArgs = pipArgs.WithXFormers("==0.0.22.post4"); + case TorchVersion.Cuda: + pipArgs = pipArgs.WithXFormers("==0.0.22.post4"); + break; + case TorchVersion.Mps: + pipArgs = pipArgs.AddArg("mpmath==1.3.0"); + break; } var requirements = new FilePath(installLocation, "requirements.txt"); diff --git a/StabilityMatrix.Core/Models/Packages/Extensions/ExtensionSpecifier.cs b/StabilityMatrix.Core/Models/Packages/Extensions/ExtensionSpecifier.cs new file mode 100644 index 00000000..13003ca0 --- /dev/null +++ b/StabilityMatrix.Core/Models/Packages/Extensions/ExtensionSpecifier.cs @@ -0,0 +1,105 @@ +using System.Diagnostics.CodeAnalysis; +using System.Text.RegularExpressions; +using JetBrains.Annotations; +using Semver; +using StabilityMatrix.Core.Processes; + +namespace StabilityMatrix.Core.Models.Packages.Extensions; + +/// +/// Extension specifier with optional version constraints. +/// +[PublicAPI] +public partial class ExtensionSpecifier +{ + public required string Name { get; init; } + + public string? Constraint { get; init; } + + public string? Version { get; init; } + + public string? VersionConstraint => Constraint is null || Version is null ? null : Constraint + Version; + + public bool TryGetSemVersionRange([NotNullWhen(true)] out SemVersionRange? semVersionRange) + { + if (!string.IsNullOrEmpty(VersionConstraint)) + { + return SemVersionRange.TryParse( + VersionConstraint, + SemVersionRangeOptions.Loose, + out semVersionRange + ); + } + + semVersionRange = null; + return false; + } + + public static ExtensionSpecifier Parse(string value) + { + TryParse(value, true, out var packageSpecifier); + + return packageSpecifier!; + } + + public static bool TryParse(string value, [NotNullWhen(true)] out ExtensionSpecifier? extensionSpecifier) + { + return TryParse(value, false, out extensionSpecifier); + } + + private static bool TryParse( + string value, + bool throwOnFailure, + [NotNullWhen(true)] out ExtensionSpecifier? packageSpecifier + ) + { + var match = ExtensionSpecifierRegex().Match(value); + if (!match.Success) + { + if (throwOnFailure) + { + throw new ArgumentException($"Invalid extension specifier: {value}"); + } + + packageSpecifier = null; + return false; + } + + packageSpecifier = new ExtensionSpecifier + { + Name = match.Groups["extension_name"].Value, + Constraint = match.Groups["version_constraint"].Value, + Version = match.Groups["version"].Value + }; + + return true; + } + + /// + public override string ToString() + { + return Name + VersionConstraint; + } + + public static implicit operator Argument(ExtensionSpecifier specifier) + { + return specifier.VersionConstraint is null + ? new Argument(specifier.Name) + : new Argument((specifier.Name, specifier.VersionConstraint)); + } + + public static implicit operator ExtensionSpecifier(string specifier) + { + return Parse(specifier); + } + + /// + /// Regex to match a pip package specifier. + /// + [GeneratedRegex( + @"(?\S+)\s*(?==|>=|<=|>|<|~=|!=)?\s*(?[a-zA-Z0-9_.]+)?", + RegexOptions.CultureInvariant, + 5000 + )] + private static partial Regex ExtensionSpecifierRegex(); +} diff --git a/StabilityMatrix.Core/Models/Packages/Extensions/GitPackageExtensionManager.cs b/StabilityMatrix.Core/Models/Packages/Extensions/GitPackageExtensionManager.cs index 1aac49f1..3cc5e17b 100644 --- a/StabilityMatrix.Core/Models/Packages/Extensions/GitPackageExtensionManager.cs +++ b/StabilityMatrix.Core/Models/Packages/Extensions/GitPackageExtensionManager.cs @@ -201,6 +201,31 @@ public abstract partial class GitPackageExtensionManager(IPrerequisiteHelper pre return extensions; } + public virtual async Task GetInstalledExtensionInfoAsync( + InstalledPackageExtension installedExtension + ) + { + if (installedExtension.PrimaryPath is not DirectoryPath extensionDirectory) + { + return installedExtension; + } + + // Get git version + var version = await prerequisiteHelper + .GetGitRepositoryVersion(extensionDirectory) + .ConfigureAwait(false); + + return installedExtension with + { + Version = new PackageExtensionVersion + { + Tag = version.Tag, + Branch = version.Branch, + CommitSha = version.CommitSha + } + }; + } + /// public virtual async Task InstallExtensionAsync( PackageExtension extension, diff --git a/StabilityMatrix.Core/Models/Packages/Extensions/IPackageExtensionManager.cs b/StabilityMatrix.Core/Models/Packages/Extensions/IPackageExtensionManager.cs index 20d37365..cdd3e3cf 100644 --- a/StabilityMatrix.Core/Models/Packages/Extensions/IPackageExtensionManager.cs +++ b/StabilityMatrix.Core/Models/Packages/Extensions/IPackageExtensionManager.cs @@ -95,6 +95,12 @@ public interface IPackageExtensionManager Task> GetInstalledExtensionsLiteAsync( InstalledPackage installedPackage, CancellationToken cancellationToken = default + + /// + /// Get updated info (version) for an installed extension. + /// + Task GetInstalledExtensionInfoAsync( + InstalledPackageExtension installedExtension ); /// diff --git a/StabilityMatrix.Core/Models/Packages/KohyaSs.cs b/StabilityMatrix.Core/Models/Packages/KohyaSs.cs index c9cdf041..ee197ac2 100644 --- a/StabilityMatrix.Core/Models/Packages/KohyaSs.cs +++ b/StabilityMatrix.Core/Models/Packages/KohyaSs.cs @@ -116,12 +116,24 @@ public class KohyaSs( Action? 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 @@ public class KohyaSs( 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,41 @@ public class KohyaSs( 1.ToPython() ); - var filesToUpdate = new[] + var kohyaGuiDir = Path.Combine(installedPackagePath, "kohya_gui"); + var guiDirExists = Directory.Exists(kohyaGuiDir); + var filesToUpdate = new List(); + + 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; diff --git a/StabilityMatrix.Core/Models/Packages/SDWebForge.cs b/StabilityMatrix.Core/Models/Packages/SDWebForge.cs index c603532c..69a501d9 100644 --- a/StabilityMatrix.Core/Models/Packages/SDWebForge.cs +++ b/StabilityMatrix.Core/Models/Packages/SDWebForge.cs @@ -39,6 +39,8 @@ public class SDWebForge( public override bool ShouldIgnoreReleases => true; public override IPackageExtensionManager ExtensionManager => null; public override string OutputFolderName => "output"; + public override PackageDifficulty InstallerSortOrder => PackageDifficulty.ReallyRecommended; + public override Dictionary>? SharedOutputFolders => new() { @@ -52,6 +54,27 @@ public class SDWebForge( public override List LaunchOptions => [ + new LaunchOptionDefinition + { + Name = "Host", + Type = LaunchOptionType.String, + DefaultValue = "localhost", + Options = ["--server-name"] + }, + new LaunchOptionDefinition + { + Name = "Port", + Type = LaunchOptionType.String, + DefaultValue = "7860", + Options = ["--port"] + }, + new LaunchOptionDefinition + { + Name = "Share", + Type = LaunchOptionType.Bool, + Description = "Set whether to share on Gradio", + Options = { "--share" } + }, new LaunchOptionDefinition { Name = "Always Offload from VRAM", @@ -155,4 +178,9 @@ public class SDWebForge( await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false); progress?.Report(new ProgressReport(1f, "Install complete", isIndeterminate: false)); } + + public override string? ExtraLaunchArguments { get; set; } = + settingsManager.IsLibraryDirSet + ? $"--gradio-allowed-path \"{settingsManager.ImagesDirectory}\"" + : string.Empty; } diff --git a/StabilityMatrix.Core/Models/Settings/Settings.cs b/StabilityMatrix.Core/Models/Settings/Settings.cs index e8556ba8..bf5e080f 100644 --- a/StabilityMatrix.Core/Models/Settings/Settings.cs +++ b/StabilityMatrix.Core/Models/Settings/Settings.cs @@ -134,6 +134,7 @@ public class Settings public Size OutputsImageSize { get; set; } = new(300, 300); public HolidayMode HolidayModeSetting { get; set; } = HolidayMode.Automatic; public bool IsWorkflowInfiniteScrollEnabled { get; set; } = true; + public bool IsOutputsTreeViewEnabled { get; set; } = true; [JsonIgnore] public bool IsHolidayModeActive => diff --git a/StabilityMatrix.Core/Models/Settings/TeachingTip.cs b/StabilityMatrix.Core/Models/Settings/TeachingTip.cs index 756e6cce..92a6ea3d 100644 --- a/StabilityMatrix.Core/Models/Settings/TeachingTip.cs +++ b/StabilityMatrix.Core/Models/Settings/TeachingTip.cs @@ -13,6 +13,7 @@ public record TeachingTip(string Value) : StringValue(Value) public static TeachingTip CheckpointCategoriesTip => new("CheckpointCategoriesTip"); public static TeachingTip PackageExtensionsInstallNotice => new("PackageExtensionsInstallNotice"); public static TeachingTip DownloadsTip => new("DownloadsTip"); + public static TeachingTip WebUiButtonMovedTip => new("WebUiButtonMovedTip"); /// public override string ToString() diff --git a/StabilityMatrix.Core/Services/DownloadService.cs b/StabilityMatrix.Core/Services/DownloadService.cs index 1aeeb9d7..ffc44ebf 100644 --- a/StabilityMatrix.Core/Services/DownloadService.cs +++ b/StabilityMatrix.Core/Services/DownloadService.cs @@ -1,4 +1,5 @@ -using System.Net.Http.Headers; +using System.Net; +using System.Net.Http.Headers; using Microsoft.Extensions.Logging; using Polly.Contrib.WaitAndRetry; using StabilityMatrix.Core.Attributes; @@ -184,6 +185,10 @@ public class DownloadService : IDownloadService throw new UnauthorizedAccessException(); } } + else if (noRedirectResponse.StatusCode == HttpStatusCode.Unauthorized) + { + throw new UnauthorizedAccessException(); + } using var redirectRequest = new HttpRequestMessage(); redirectRequest.Method = HttpMethod.Get; diff --git a/StabilityMatrix.Tests/Helper/PackageFactoryTests.cs b/StabilityMatrix.Tests/Helper/PackageFactoryTests.cs index 4163922d..b4920bdb 100644 --- a/StabilityMatrix.Tests/Helper/PackageFactoryTests.cs +++ b/StabilityMatrix.Tests/Helper/PackageFactoryTests.cs @@ -8,31 +8,28 @@ public class PackageFactoryTests { private PackageFactory packageFactory = null!; private IEnumerable fakeBasePackages = null!; - + [TestInitialize] public void Setup() { - fakeBasePackages = new List - { - new DankDiffusion(null!, null!, null!, null!) - }; - packageFactory = new PackageFactory(fakeBasePackages); + fakeBasePackages = new List { new DankDiffusion(null!, null!, null!, null!) }; + packageFactory = new PackageFactory(fakeBasePackages, null, null, null, null, null); } - + [TestMethod] public void GetAllAvailablePackages_ReturnsAllPackages() { var result = packageFactory.GetAllAvailablePackages(); Assert.AreEqual(1, result.Count()); } - + [TestMethod] public void FindPackageByName_ReturnsPackage() { var result = packageFactory.FindPackageByName("dank-diffusion"); Assert.IsNotNull(result); } - + [TestMethod] public void FindPackageByName_ReturnsNull() {