Merge pull request LykosAI#750 from ionite34/disposables-improvement

Add IDisposable return for RelayPropertyFor subscriptions, and use dispose for transients

Author: ionite34

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 ### Fixed
 - Fixed "The version of the native libSkiaSharp library (88.1) is incompatible with this version of SkiaSharp." error for Linux users
 - Fixed download links for IPAdapters in the HuggingFace model browser
+- Fixed potential memory leak of transient controls (Inference Prompt and Output Image Viewer) not being garbage collected due to event subscriptions
 
 ## v2.12.0-dev.2
 ### Added
@@ -58,6 +59,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 ## v2.11.6
 ### Fixed
 - Fixed incorrect IPAdapter download links in the HuggingFace model browser
+- Fixed potential memory leak of transient controls (Inference Prompt and Output Image Viewer) not being garbage collected due to event subscriptions
 
 ## v2.11.5
 ### Added

StabilityMatrix.Avalonia/ViewModels/Base/DisposableLoadableViewModelBase.cs

@@ -0,0 +1,75 @@
+using System;
+using System.Reactive.Disposables;
+using JetBrains.Annotations;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Base;
+
+public abstract class DisposableLoadableViewModelBase : LoadableViewModelBase, IDisposable
+{
+    private readonly CompositeDisposable instanceDisposables = new();
+
+    /// <summary>
+    /// Adds a disposable to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposable">The disposable to add.</param>
+    protected void AddDisposable([HandlesResourceDisposal] IDisposable disposable)
+    {
+        instanceDisposables.Add(disposable);
+    }
+
+    /// <summary>
+    /// Adds disposables to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposables">The disposables to add.</param>
+    protected void AddDisposable([HandlesResourceDisposal] params IDisposable[] disposables)
+    {
+        foreach (var disposable in disposables)
+        {
+            instanceDisposables.Add(disposable);
+        }
+    }
+
+    /// <summary>
+    /// Adds a disposable to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposable">The disposable to add.</param>
+    /// <typeparam name="T">The type of the disposable.</typeparam>
+    /// <returns>The disposable that was added.</returns>
+    protected T AddDisposable<T>([HandlesResourceDisposal] T disposable)
+        where T : IDisposable
+    {
+        instanceDisposables.Add(disposable);
+        return disposable;
+    }
+
+    /// <summary>
+    /// Adds disposables to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposables">The disposables to add.</param>
+    /// <typeparam name="T">The type of the disposables.</typeparam>
+    /// <returns>The disposables that were added.</returns>
+    protected T[] AddDisposable<T>([HandlesResourceDisposal] params T[] disposables)
+        where T : IDisposable
+    {
+        foreach (var disposable in disposables)
+        {
+            instanceDisposables.Add(disposable);
+        }
+
+        return disposables;
+    }
+
+    protected virtual void Dispose(bool disposing)
+    {
+        if (disposing)
+        {
+            instanceDisposables.Dispose();
+        }
+    }
+
+    public void Dispose()
+    {
+        Dispose(true);
+        GC.SuppressFinalize(this);
+    }
+}

StabilityMatrix.Avalonia/ViewModels/Base/DisposableViewModelBase.cs

@@ -0,0 +1,75 @@
+using System;
+using System.Reactive.Disposables;
+using JetBrains.Annotations;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Base;
+
+public abstract class DisposableViewModelBase : ViewModelBase, IDisposable
+{
+    private readonly CompositeDisposable instanceDisposables = new();
+
+    /// <summary>
+    /// Adds a disposable to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposable">The disposable to add.</param>
+    protected void AddDisposable([HandlesResourceDisposal] IDisposable disposable)
+    {
+        instanceDisposables.Add(disposable);
+    }
+
+    /// <summary>
+    /// Adds disposables to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposables">The disposables to add.</param>
+    protected void AddDisposable([HandlesResourceDisposal] params IDisposable[] disposables)
+    {
+        foreach (var disposable in disposables)
+        {
+            instanceDisposables.Add(disposable);
+        }
+    }
+
+    /// <summary>
+    /// Adds a disposable to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposable">The disposable to add.</param>
+    /// <typeparam name="T">The type of the disposable.</typeparam>
+    /// <returns>The disposable that was added.</returns>
+    protected T AddDisposable<T>([HandlesResourceDisposal] T disposable)
+        where T : IDisposable
+    {
+        instanceDisposables.Add(disposable);
+        return disposable;
+    }
+
+    /// <summary>
+    /// Adds disposables to be disposed when this view model is disposed.
+    /// </summary>
+    /// <param name="disposables">The disposables to add.</param>
+    /// <typeparam name="T">The type of the disposables.</typeparam>
+    /// <returns>The disposables that were added.</returns>
+    protected T[] AddDisposable<T>([HandlesResourceDisposal] params T[] disposables)
+        where T : IDisposable
+    {
+        foreach (var disposable in disposables)
+        {
+            instanceDisposables.Add(disposable);
+        }
+
+        return disposables;
+    }
+
+    protected virtual void Dispose(bool disposing)
+    {
+        if (disposing)
+        {
+            instanceDisposables.Dispose();
+        }
+    }
+
+    public void Dispose()
+    {
+        Dispose(true);
+        GC.SuppressFinalize(this);
+    }
+}

StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs

@@ -93,7 +93,7 @@ RunningPackageService runningPackageService
         ClientManager = inferenceClientManager;
 
         ImageGalleryCardViewModel = vmFactory.Get<ImageGalleryCardViewModel>();
-        ImageFolderCardViewModel = vmFactory.Get<ImageFolderCardViewModel>();
+        ImageFolderCardViewModel = AddDisposable(vmFactory.Get<ImageFolderCardViewModel>());
 
         GenerateImageCommand.WithConditionalNotificationErrorHandler(notificationService);
     }

StabilityMatrix.Avalonia/ViewModels/Base/InferenceTabViewModelBase.cs

@@ -33,8 +33,7 @@
 namespace StabilityMatrix.Avalonia.ViewModels.Base;
 
 public abstract partial class InferenceTabViewModelBase
-    : LoadableViewModelBase,
-        IDisposable,
+    : DisposableLoadableViewModelBase,
         IPersistentViewProvider,
         IDropTarget
 {
@@ -158,21 +157,16 @@ private async Task DebugLoadViewState()
         }
     }
 
-    protected virtual void Dispose(bool disposing)
+    protected override void Dispose(bool disposing)
     {
+        base.Dispose(disposing);
+
         if (disposing)
         {
             ((IPersistentViewProvider)this).AttachedPersistentView = null;
         }
     }
 
-    /// <inheritdoc />
-    public void Dispose()
-    {
-        Dispose(true);
-        GC.SuppressFinalize(this);
-    }
-
     /// <summary>
     /// Loads image and metadata from a file path
     /// </summary>

StabilityMatrix.Avalonia/ViewModels/Inference/ImageFolderCardViewModel.cs

@@ -36,7 +36,7 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 [View(typeof(ImageFolderCard))]
 [ManagedService]
 [Transient]
-public partial class ImageFolderCardViewModel : ViewModelBase
+public partial class ImageFolderCardViewModel : DisposableViewModelBase
 {
     private readonly ILogger<ImageFolderCardViewModel> logger;
     private readonly IImageIndexService imageIndexService;
@@ -86,11 +86,13 @@ ServiceManager<ViewModelBase> vmFactory
             .Bind(LocalImages)
             .Subscribe();
 
-        settingsManager.RelayPropertyFor(
-            this,
-            vm => vm.ImageSize,
-            settings => settings.InferenceImageSize,
-            delay: TimeSpan.FromMilliseconds(250)
+        AddDisposable(
+            settingsManager.RelayPropertyFor(
+                this,
+                vm => vm.ImageSize,
+                settings => settings.InferenceImageSize,
+                delay: TimeSpan.FromMilliseconds(250)
+            )
         );
     }
 

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -83,7 +83,8 @@ RunningPackageService runningPackageService
             samplerCard.DenoiseStrength = 1.0d;
         });
 
-        PromptCardViewModel = vmFactory.Get<PromptCardViewModel>();
+        PromptCardViewModel = AddDisposable(vmFactory.Get<PromptCardViewModel>());
+
         BatchSizeCardViewModel = vmFactory.Get<BatchSizeCardViewModel>();
 
         ModulesCardViewModel = vmFactory.Get<StackEditableCardViewModel>(modulesCard =>
@@ -109,13 +110,15 @@ RunningPackageService runningPackageService
         );
 
         // When refiner is provided in model card, enable for sampler
-        ModelCardViewModel
-            .WhenPropertyChanged(x => x.IsRefinerSelectionEnabled)
-            .Subscribe(e =>
-            {
-                SamplerCardViewModel.IsRefinerStepsEnabled =
-                    e.Sender is { IsRefinerSelectionEnabled: true, SelectedRefiner: not null };
-            });
+        AddDisposable(
+            ModelCardViewModel
+                .WhenPropertyChanged(x => x.IsRefinerSelectionEnabled)
+                .Subscribe(e =>
+                {
+                    SamplerCardViewModel.IsRefinerStepsEnabled =
+                        e.Sender is { IsRefinerSelectionEnabled: true, SelectedRefiner: not null };
+                })
+        );
     }
 
     /// <inheritdoc />

StabilityMatrix.Avalonia/ViewModels/Inference/PromptCardViewModel.cs

@@ -28,7 +28,10 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 [View(typeof(PromptCard))]
 [ManagedService]
 [Transient]
-public partial class PromptCardViewModel : LoadableViewModelBase, IParametersLoadableState, IComfyStep
+public partial class PromptCardViewModel
+    : DisposableLoadableViewModelBase,
+        IParametersLoadableState,
+        IComfyStep
 {
     private readonly IModelIndexService modelIndexService;
     private readonly ISettingsManager settingsManager;
@@ -75,11 +78,13 @@ SharedState sharedState
             vm.AvailableModules = [typeof(PromptExpansionModule)];
         });
 
-        settingsManager.RelayPropertyFor(
-            this,
-            vm => vm.IsAutoCompletionEnabled,
-            settings => settings.IsPromptCompletionEnabled,
-            true
+        AddDisposable(
+            settingsManager.RelayPropertyFor(
+                this,
+                vm => vm.IsAutoCompletionEnabled,
+                settings => settings.IsPromptCompletionEnabled,
+                true
+            )
         );
     }
 

StabilityMatrix.Core/Services/ISettingsManager.cs

@@ -73,7 +73,7 @@ public interface ISettingsManager
     /// <summary>
     /// Register a source observable object and property to be relayed to Settings
     /// </summary>
-    void RelayPropertyFor<T, TValue>(
+    IDisposable RelayPropertyFor<T, TValue>(
         T source,
         Expression<Func<T, TValue>> sourceProperty,
         Expression<Func<Settings, TValue>> settingsProperty,

StabilityMatrix.Core/Services/SettingsManager.cs

@@ -1,6 +1,7 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq.Expressions;
+using System.Reactive.Disposables;
 using System.Reactive.Linq;
 using System.Reflection;
 using System.Text.Json;
@@ -163,7 +164,7 @@ public void Transaction<TValue>(Expression<Func<Settings, TValue>> expression, T
     }
 
     /// <inheritdoc />
-    public void RelayPropertyFor<T, TValue>(
+    public IDisposable RelayPropertyFor<T, TValue>(
         T source,
         Expression<Func<T, TValue>> sourceProperty,
         Expression<Func<Settings, TValue>> settingsProperty,
@@ -181,7 +182,7 @@ public void RelayPropertyFor<T, TValue>(
         var sourceTypeName = source.GetType().Name;
 
         // Update source when settings change
-        SettingsPropertyChanged += (sender, args) =>
+        void OnSettingsPropertyChanged(object? sender, RelayPropertyChangedEventArgs args)
         {
             if (args.PropertyName != settingsPropertyPath)
                 return;
@@ -198,10 +199,10 @@ public void RelayPropertyFor<T, TValue>(
             );
 
             sourceInstanceAccessor.Set(source, settingsAccessor.Get(Settings));
-        };
+        }
 
         // Set and Save settings when source changes
-        source.PropertyChanged += (sender, args) =>
+        void OnSourcePropertyChanged(object? sender, PropertyChangedEventArgs args)
         {
             if (args.PropertyName != sourcePropertyPath)
                 return;
@@ -241,13 +242,32 @@ public void RelayPropertyFor<T, TValue>(
                 sender,
                 new RelayPropertyChangedEventArgs(settingsPropertyPath, true)
             );
-        };
+        }
+
+        var subscription = Disposable.Create(() =>
+        {
+            source.PropertyChanged -= OnSourcePropertyChanged;
+            SettingsPropertyChanged -= OnSettingsPropertyChanged;
+        });
 
-        // Set initial value if requested
-        if (setInitial)
+        try
         {
-            sourceInstanceAccessor.Set(settingsAccessor.Get(Settings));
+            SettingsPropertyChanged += OnSettingsPropertyChanged;
+            source.PropertyChanged += OnSourcePropertyChanged;
+
+            // Set initial value if requested
+            if (setInitial)
+            {
+                sourceInstanceAccessor.Set(settingsAccessor.Get(Settings));
+            }
         }
+        catch
+        {
+            subscription.Dispose();
+            throw;
+        }
+
+        return subscription;
     }
 
     /// <inheritdoc />