Merge pull request #312 from ionite34/auto-retry-ioexceptions

Added auto-retry for any IOException that happens during a download

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - NVIDIA GPU users will be updated to use CUDA 12.1 for ComfyUI & Fooocus packages for a slight performance improvement
   - Update will occur the next time the package is updated, or on a fresh install
   - Note: CUDA 12.1 is only available on Maxwell (GTX 900 series) and newer GPUs
+- Improved Model Browser download stability with automatic retries for download errors
 ### Fixed
 - Fixed crash when clicking Inference gallery image after the image is deleted externally in file explorer
 - Fixed Inference popup Install button not working on One-Click Installer

StabilityMatrix.Core/Models/TrackedDownload.cs

@@ -68,6 +68,8 @@ public class TrackedDownload
     [JsonIgnore]
     public Exception? Exception { get; private set; }
 
+    private int attempts;
+
     #region Events
     private WeakEventManager<ProgressReport>? progressUpdateEventManager;
 
@@ -120,22 +122,14 @@ public class TrackedDownload
     private async Task StartDownloadTask(long resumeFromByte, CancellationToken cancellationToken)
     {
         var progress = new Progress<ProgressReport>(OnProgressUpdate);
-        try
-        {
-            await downloadService!
-                .ResumeDownloadToFileAsync(
-                    SourceUrl.ToString(),
-                    DownloadDirectory.JoinFile(TempFileName),
-                    resumeFromByte,
-                    progress,
-                    cancellationToken: cancellationToken
-                )
-                .ConfigureAwait(false);
-        }
-        catch (Exception e)
-        {
-            Logger.Warn(e);
-        }
+
+        await downloadService!.ResumeDownloadToFileAsync(
+            SourceUrl.ToString(),
+            DownloadDirectory.JoinFile(TempFileName),
+            resumeFromByte,
+            progress,
+            cancellationToken: cancellationToken
+        );
 
         // If hash validation is enabled, validate the hash
         if (ValidateHash)
@@ -313,6 +307,23 @@ public class TrackedDownload
             // Set the exception
             Exception = task.Exception;
 
+            if (
+                (Exception is IOException || Exception?.InnerException is IOException)
+                && attempts < 3
+            )
+            {
+                attempts++;
+                Logger.Warn(
+                    "Download {Download} failed with {Exception}, retrying ({Attempt})",
+                    FileName,
+                    Exception,
+                    attempts
+                );
+                ProgressState = ProgressState.Inactive;
+                Resume();
+                return;
+            }
+
             ProgressState = ProgressState.Failed;
         }
         // Otherwise success