From a5beaf434ced3b8240fb6e984a3519d9c660b628 Mon Sep 17 00:00:00 2001 From: JT Date: Thu, 14 Mar 2024 22:06:38 -0700 Subject: [PATCH] handle download links from sm chrome extension & fix model delete not deleting from cache --- CHANGELOG.md | 3 + .../Services/IModelDownloadLinkHandler.cs | 8 + .../Services/ModelDownloadLinkHandler.cs | 248 ++++++++++++++++++ .../CivitAiBrowserViewModel.cs | 5 +- .../CheckpointManager/CheckpointFile.cs | 7 + 5 files changed, 270 insertions(+), 1 deletion(-) create mode 100644 StabilityMatrix.Avalonia/Services/IModelDownloadLinkHandler.cs create mode 100644 StabilityMatrix.Avalonia/Services/ModelDownloadLinkHandler.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index 91a94ec3..0a543042 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ 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.3 +### Added +- Added support for deep links from the new Stability Matrix Chrome extension ### 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 @@ -13,6 +15,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2 - 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 ## v2.10.0-dev.2 ### Added diff --git a/StabilityMatrix.Avalonia/Services/IModelDownloadLinkHandler.cs b/StabilityMatrix.Avalonia/Services/IModelDownloadLinkHandler.cs new file mode 100644 index 00000000..dfea3c3f --- /dev/null +++ b/StabilityMatrix.Avalonia/Services/IModelDownloadLinkHandler.cs @@ -0,0 +1,8 @@ +using System.Threading.Tasks; + +namespace StabilityMatrix.Avalonia.Services; + +public interface IModelDownloadLinkHandler +{ + Task StartListening(); +} diff --git a/StabilityMatrix.Avalonia/Services/ModelDownloadLinkHandler.cs b/StabilityMatrix.Avalonia/Services/ModelDownloadLinkHandler.cs new file mode 100644 index 00000000..326d715c --- /dev/null +++ b/StabilityMatrix.Avalonia/Services/ModelDownloadLinkHandler.cs @@ -0,0 +1,248 @@ +using System; +using System.IO; +using System.Linq; +using System.Threading.Tasks; +using System.Web; +using Avalonia.Controls.Notifications; +using Avalonia.Threading; +using MessagePipe; +using Microsoft.Extensions.Logging; +using StabilityMatrix.Avalonia.Helpers; +using StabilityMatrix.Core.Api; +using StabilityMatrix.Core.Attributes; +using StabilityMatrix.Core.Extensions; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api; +using StabilityMatrix.Core.Models.FileInterfaces; +using StabilityMatrix.Core.Services; + +namespace StabilityMatrix.Avalonia.Services; + +[Singleton(typeof(IModelDownloadLinkHandler)), Singleton(typeof(IAsyncDisposable))] +public class ModelDownloadLinkHandler( + IDistributedSubscriber uriHandlerSubscriber, + ILogger logger, + ICivitApi civitApi, + INotificationService notificationService, + ISettingsManager settingsManager, + IDownloadService downloadService, + ITrackedDownloadService trackedDownloadService +) : IAsyncDisposable, IModelDownloadLinkHandler +{ + private IAsyncDisposable? uriHandlerSubscription; + private const string DownloadCivitModel = "downloadCivitModel"; + + public async Task StartListening() + { + uriHandlerSubscription = await uriHandlerSubscriber.SubscribeAsync( + UriHandler.IpcKeySend, + UriReceivedHandler + ); + } + + public async ValueTask DisposeAsync() + { + if (uriHandlerSubscription is not null) + { + await uriHandlerSubscription.DisposeAsync(); + uriHandlerSubscription = null; + } + } + + private void UriReceivedHandler(Uri receivedUri) + { + logger.LogDebug("ModelDownloadLinkHandler Received URI: {Uri}", receivedUri.PathAndQuery); + if (!receivedUri.Host.Equals(DownloadCivitModel, StringComparison.OrdinalIgnoreCase)) + return; + + var queryDict = HttpUtility.ParseQueryString(receivedUri.Query); + var modelIdStr = queryDict["modelId"]; + var modelVersionIdStr = queryDict["modelVersionId"]; + var type = queryDict["type"]; + var format = queryDict["format"]; + var size = queryDict["size"]; + var fp = queryDict["fp"]; + + if ( + string.IsNullOrWhiteSpace(modelIdStr) + || string.IsNullOrWhiteSpace(type) + || string.IsNullOrWhiteSpace(format) + || !int.TryParse(modelIdStr, out var modelId) + || !Enum.TryParse(type, out var civitFileType) + || !Enum.TryParse(format, out var civitFormat) + ) + { + logger.LogError("ModelDownloadLinkHandler: Invalid query parameters"); + + Dispatcher.UIThread.Post( + () => + notificationService.Show( + new Notification( + "Invalid Download Link", + "The download link is invalid", + NotificationType.Error + ) + ) + ); + return; + } + + Dispatcher.UIThread.Post( + () => + notificationService.Show( + "Link Received", + "Successfully received download link", + NotificationType.Warning + ) + ); + + var modelTask = civitApi.GetModelById(modelId); + modelTask.Wait(); + var model = modelTask.Result; + + var useModelVersion = !string.IsNullOrWhiteSpace(modelVersionIdStr); + var modelVersionId = useModelVersion ? int.Parse(modelVersionIdStr) : 0; + + var modelVersion = useModelVersion + ? model.ModelVersions?.FirstOrDefault(x => x.Id == modelVersionId) + : model.ModelVersions?.FirstOrDefault(); + + if (modelVersion is null) + { + logger.LogError("ModelDownloadLinkHandler: Model version not found"); + Dispatcher.UIThread.Post( + () => + notificationService.Show( + new Notification( + "Model has no versions available", + "This model has no versions available for download", + NotificationType.Error + ) + ) + ); + return; + } + + var possibleFiles = modelVersion.Files?.Where( + x => x.Type == civitFileType && x.Metadata.Format == civitFormat + ); + + if (!string.IsNullOrWhiteSpace(fp) && Enum.TryParse(fp, out var fpType)) + { + possibleFiles = possibleFiles?.Where(x => x.Metadata.Fp == fpType); + } + + if (!string.IsNullOrWhiteSpace(size) && Enum.TryParse(size, out var modelSize)) + { + possibleFiles = possibleFiles?.Where(x => x.Metadata.Size == modelSize); + } + + possibleFiles = possibleFiles?.ToList(); + + if (possibleFiles is null) + { + Dispatcher.UIThread.Post( + () => + notificationService.Show( + new Notification( + "Model has no files available", + "This model has no files available for download", + NotificationType.Error + ) + ) + ); + logger.LogError("ModelDownloadLinkHandler: Model file not found"); + return; + } + + var selectedFile = possibleFiles.FirstOrDefault() ?? modelVersion.Files?.FirstOrDefault(); + + var rootModelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory); + var downloadDirectory = rootModelsDirectory.JoinDir( + selectedFile.Type == CivitFileType.VAE + ? SharedFolderType.VAE.GetStringValue() + : model.Type.ConvertTo().GetStringValue() + ); + + downloadDirectory.Create(); + var downloadPath = downloadDirectory.JoinFile(selectedFile.Name); + + // Create tracked download + var download = trackedDownloadService.NewDownload(selectedFile.DownloadUrl, downloadPath); + + // Download model info and preview first + var saveCmInfoTask = SaveCmInfo(model, modelVersion, selectedFile, downloadDirectory); + var savePreviewImageTask = SavePreviewImage(modelVersion, downloadPath); + + Task.WaitAll([saveCmInfoTask, savePreviewImageTask]); + + var cmInfoPath = saveCmInfoTask.Result; + var previewImagePath = savePreviewImageTask.Result; + + // Add hash info + download.ExpectedHashSha256 = selectedFile.Hashes.SHA256; + + // Add files to cleanup list + download.ExtraCleanupFileNames.Add(cmInfoPath); + if (previewImagePath is not null) + { + download.ExtraCleanupFileNames.Add(previewImagePath); + } + + // Add hash context action + download.ContextAction = CivitPostDownloadContextAction.FromCivitFile(selectedFile); + + download.Start(); + + Dispatcher.UIThread.Post( + () => notificationService.Show("Download Started", $"Downloading {selectedFile.Name}") + ); + } + + private static async Task SaveCmInfo( + CivitModel model, + CivitModelVersion modelVersion, + CivitFile modelFile, + DirectoryPath downloadDirectory + ) + { + var modelFileName = Path.GetFileNameWithoutExtension(modelFile.Name); + var modelInfo = new ConnectedModelInfo(model, modelVersion, modelFile, DateTime.UtcNow); + + await modelInfo.SaveJsonToDirectory(downloadDirectory, modelFileName); + + var jsonName = $"{modelFileName}.cm-info.json"; + return downloadDirectory.JoinFile(jsonName); + } + + /// + /// Saves the preview image to the same directory as the model file + /// + /// + /// + /// The file path of the saved preview image + private async Task SavePreviewImage(CivitModelVersion modelVersion, FilePath modelFilePath) + { + // Skip if model has no images + if (modelVersion.Images == null || modelVersion.Images.Count == 0) + { + return null; + } + + var image = modelVersion.Images[0]; + var imageExtension = Path.GetExtension(image.Url).TrimStart('.'); + if (imageExtension is "jpg" or "jpeg" or "png") + { + var imageDownloadPath = modelFilePath.Directory!.JoinFile( + $"{modelFilePath.NameWithoutExtension}.preview.{imageExtension}" + ); + + var imageTask = downloadService.DownloadToFileAsync(image.Url, imageDownloadPath); + await notificationService.TryAsync(imageTask, "Could not download preview image"); + + return imageDownloadPath; + } + + return null; + } +} diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs index 924d17ed..a1ffd4e4 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivitAiBrowserViewModel.cs @@ -139,7 +139,8 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase ISettingsManager settingsManager, ServiceManager dialogFactory, ILiteDbContext liteDbContext, - INotificationService notificationService + INotificationService notificationService, + IModelDownloadLinkHandler modelDownloadLinkHandler ) { this.civitApi = civitApi; @@ -163,6 +164,8 @@ public partial class CivitAiBrowserViewModel : TabViewModelBase .Subscribe(_ => TrySearchAgain(false).SafeFireAndForget(), err => Logger.Error(err)); EventManager.Instance.NavigateAndFindCivitModelRequested += OnNavigateAndFindCivitModelRequested; + + modelDownloadLinkHandler.StartListening().SafeFireAndForget(); } private void OnNavigateAndFindCivitModelRequested(object? sender, int e) diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFile.cs index c7398817..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; @@ -136,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)