From 80d1b86ecb8a5da0bdda6705dc50253b78195f1d Mon Sep 17 00:00:00 2001 From: Ionite Date: Mon, 3 Jul 2023 17:14:29 -0400 Subject: [PATCH] Improve checkpoint browser download error handling --- StabilityMatrix/Models/ConnectedModelInfo.cs | 5 +- .../Models/FileInterfaces/FileSystemPath.cs | 6 +- .../CheckpointBrowserCardViewModel.cs | 190 +++++++++++------- 3 files changed, 122 insertions(+), 79 deletions(-) diff --git a/StabilityMatrix/Models/ConnectedModelInfo.cs b/StabilityMatrix/Models/ConnectedModelInfo.cs index a881c95d..12a44bfe 100644 --- a/StabilityMatrix/Models/ConnectedModelInfo.cs +++ b/StabilityMatrix/Models/ConnectedModelInfo.cs @@ -9,6 +9,9 @@ namespace StabilityMatrix.Models; public class ConnectedModelInfo { + [JsonIgnore] + public const string FileExtension = ".cm-info.json"; + public int ModelId { get; set; } public string ModelName { get; set; } public string ModelDescription { get; set; } @@ -61,7 +64,7 @@ public class ConnectedModelInfo /// Model file name without extensions public async Task SaveJsonToDirectory(string directoryPath, string modelFileName) { - var name = modelFileName + ".cm-info.json"; + var name = modelFileName + FileExtension; var json = JsonSerializer.Serialize(this); await File.WriteAllTextAsync(Path.Combine(directoryPath, name), json); } diff --git a/StabilityMatrix/Models/FileInterfaces/FileSystemPath.cs b/StabilityMatrix/Models/FileInterfaces/FileSystemPath.cs index 783c2651..272ae6cb 100644 --- a/StabilityMatrix/Models/FileInterfaces/FileSystemPath.cs +++ b/StabilityMatrix/Models/FileInterfaces/FileSystemPath.cs @@ -19,6 +19,11 @@ public class FileSystemPath : IEquatable, IEquatable protected FileSystemPath(params string[] paths) : this(Path.Combine(paths)) { } + + public override string ToString() + { + return FullPath; + } public bool Equals(FileSystemPath? other) { @@ -29,7 +34,6 @@ public class FileSystemPath : IEquatable, IEquatable public bool Equals(string? other) { - if (ReferenceEquals(null, other)) return false; return string.Equals(FullPath, other); } diff --git a/StabilityMatrix/ViewModels/CheckpointBrowserCardViewModel.cs b/StabilityMatrix/ViewModels/CheckpointBrowserCardViewModel.cs index c72d80d1..465e7690 100644 --- a/StabilityMatrix/ViewModels/CheckpointBrowserCardViewModel.cs +++ b/StabilityMatrix/ViewModels/CheckpointBrowserCardViewModel.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Generic; using System.IO; using System.Linq; using System.Net.Http; @@ -13,6 +14,7 @@ using StabilityMatrix.Extensions; using StabilityMatrix.Helper; using StabilityMatrix.Models; using StabilityMatrix.Models.Api; +using StabilityMatrix.Models.FileInterfaces; using StabilityMatrix.Models.Progress; using StabilityMatrix.Services; using Wpf.Ui.Controls; @@ -76,95 +78,129 @@ public partial class CheckpointBrowserCardViewModel : ProgressViewModel { Text = "Downloading..."; - // Get latest version - if (!(model.ModelVersions?.Count > 0)) - { - snackbarService.ShowSnackbarAsync( - "This model has no versions available for download", - "Model has no versions available", ControlAppearance.Caution).SafeFireAndForget(); - return; - } - var modelVersion = model.ModelVersions[0]; + // Holds files to be deleted on errors + var filesForCleanup = new HashSet(); - // Get latest version file - if (!(modelVersion.Files?.Count > 0)) + // Set Text when exiting, finally block will set 100 and delay clear progress + try { - snackbarService.ShowSnackbarAsync( - "This model has no files available for download", - "Model has no files available", ControlAppearance.Caution).SafeFireAndForget(); - return; - } - var modelFile = modelVersion.Files[0]; - - var downloadFolder = Path.Combine(settingsManager.ModelsDirectory, - model.Type.ConvertTo().GetStringValue()); - // Folders might be missing if user didn't install any packages yet - Directory.CreateDirectory(downloadFolder); - var downloadPath = Path.GetFullPath(Path.Combine(downloadFolder, modelFile.Name)); - - // Do the download - var downloadTask = downloadService.DownloadToFileAsync(modelFile.DownloadUrl, downloadPath, - new Progress(report => + // Get latest version + if (!(model.ModelVersions?.Count > 0)) { - Value = report.Percentage; - Text = $"Downloading... {report.Percentage}%"; - })); - var downloadResult = await snackbarService.TryAsync(downloadTask, "Could not download file"); - - // For exceptions other than ApiException or TaskCanceledException, log error - if (downloadResult is {Exception: not null and not - (HttpRequestException or ApiException or TaskCanceledException)}) - { - Logger.Error(downloadResult.Exception, "Unexpected error during model download"); - return; - } - - // When sha256 is available, validate the downloaded file - var fileExpectedSha256 = modelFile.Hashes.SHA256; - if (!string.IsNullOrEmpty(fileExpectedSha256)) - { - var hashProgress = new Progress(progress => - { - Value = progress.Percentage; - Text = $"Validating... {progress.Percentage}%"; - }); - var sha256 = await FileHash.GetSha256Async(downloadPath, hashProgress); - if (sha256 != fileExpectedSha256.ToLowerInvariant()) + snackbarService.ShowSnackbarAsync( + "This model has no versions available for download", + "Model has no versions available", ControlAppearance.Caution).SafeFireAndForget(); + Text = "Unable to Download"; + return; + } + var modelVersion = model.ModelVersions[0]; + + // Get latest version file + if (!(modelVersion.Files?.Count > 0)) { - Text = "Import Failed!"; - DelayedClearProgress(TimeSpan.FromMilliseconds(800)); snackbarService.ShowSnackbarAsync( - "This may be caused by network or server issues from CivitAI, please try again in a few minutes.", - "Download failed hash validation").SafeFireAndForget(); + "This model has no files available for download", + "Model has no files available", ControlAppearance.Caution).SafeFireAndForget(); + Text = "Unable to Download"; + return; + } + var modelFile = modelVersion.Files[0]; + + var downloadFolder = Path.Combine(settingsManager.ModelsDirectory, + model.Type.ConvertTo().GetStringValue()); + // Folders might be missing if user didn't install any packages yet + Directory.CreateDirectory(downloadFolder); + var downloadPath = Path.GetFullPath(Path.Combine(downloadFolder, modelFile.Name)); + filesForCleanup.Add(downloadPath); + + // Do the download + var downloadTask = downloadService.DownloadToFileAsync(modelFile.DownloadUrl, downloadPath, + new Progress(report => + { + Value = report.Percentage; + Text = $"Downloading... {report.Percentage}%"; + })); + var downloadResult = await snackbarService.TryAsync(downloadTask, "Could not download file"); + + // Failed download handling + if (downloadResult.Exception is not null) + { + // For exceptions other than ApiException or TaskCanceledException, log error + var logLevel = downloadResult.Exception switch + { + HttpRequestException or ApiException or TaskCanceledException => LogLevel.Warn, + _ => LogLevel.Error + }; + Logger.Log(logLevel, downloadResult.Exception, "Error during model download"); + + Text = "Download Failed"; return; } - snackbarService.ShowSnackbarAsync($"{model.Type} {model.Name} imported successfully!", - "Import complete", ControlAppearance.Info).SafeFireAndForget(); + + // When sha256 is available, validate the downloaded file + var fileExpectedSha256 = modelFile.Hashes.SHA256; + if (!string.IsNullOrEmpty(fileExpectedSha256)) + { + var hashProgress = new Progress(progress => + { + Value = progress.Percentage; + Text = $"Validating... {progress.Percentage}%"; + }); + var sha256 = await FileHash.GetSha256Async(downloadPath, hashProgress); + if (sha256 != fileExpectedSha256.ToLowerInvariant()) + { + Text = "Import Failed!"; + DelayedClearProgress(TimeSpan.FromMilliseconds(800)); + snackbarService.ShowSnackbarAsync( + "This may be caused by network or server issues from CivitAI, please try again in a few minutes.", + "Download failed hash validation").SafeFireAndForget(); + Text = "Download Failed"; + return; + } + snackbarService.ShowSnackbarAsync($"{model.Type} {model.Name} imported successfully!", + "Import complete", ControlAppearance.Info).SafeFireAndForget(); + } + + IsIndeterminate = true; + + // Save connected model info + var modelFileName = Path.GetFileNameWithoutExtension(modelFile.Name); + var modelInfo = new ConnectedModelInfo(CivitModel, modelVersion, modelFile, DateTime.UtcNow); + var modelInfoPath = Path.GetFullPath(Path.Combine( + downloadFolder, modelFileName + ConnectedModelInfo.FileExtension)); + filesForCleanup.Add(modelInfoPath); + await modelInfo.SaveJsonToDirectory(downloadFolder, modelFileName); + + // If available, save a model image + if (modelVersion.Images != null && modelVersion.Images.Any()) + { + var image = modelVersion.Images[0]; + var imageExtension = Path.GetExtension(image.Url).TrimStart('.'); + if (imageExtension is "jpg" or "jpeg" or "png") + { + var imageDownloadPath = Path.GetFullPath(Path.Combine(downloadFolder, $"{modelFileName}.preview.{imageExtension}")); + filesForCleanup.Add(imageDownloadPath); + var imageTask = downloadService.DownloadToFileAsync(image.Url, imageDownloadPath); + var imageResult = await snackbarService.TryAsync(imageTask, "Could not download preview image"); + } + } + + // Successful - clear cleanup list + filesForCleanup.Clear(); + + Text = "Import complete!"; } - - IsIndeterminate = true; - - // Save connected model info - var modelFileName = Path.GetFileNameWithoutExtension(modelFile.Name); - var modelInfo = new ConnectedModelInfo(CivitModel, modelVersion, modelFile, DateTime.UtcNow); - await modelInfo.SaveJsonToDirectory(downloadFolder, modelFileName); - - // If available, save a model image - if (modelVersion.Images != null && modelVersion.Images.Any()) + finally { - var image = modelVersion.Images[0]; - var imageExtension = Path.GetExtension(image.Url).TrimStart('.'); - if (imageExtension is "jpg" or "jpeg" or "png") + foreach (var file in filesForCleanup.Where(file => file.Exists)) { - var imageDownloadPath = Path.GetFullPath(Path.Combine(downloadFolder, $"{modelFileName}.preview.{imageExtension}")); - await downloadService.DownloadToFileAsync(image.Url, imageDownloadPath); + file.Delete(); + Logger.Info($"Download cleanup: Deleted file {file}"); } + IsIndeterminate = false; + Value = 100; + DelayedClearProgress(TimeSpan.FromMilliseconds(800)); } - - IsIndeterminate = false; - Text = "Import complete!"; - Value = 100; - DelayedClearProgress(TimeSpan.FromMilliseconds(800)); } private void DelayedClearProgress(TimeSpan delay)