Add Load/Unload methods to IStableDiffusionService

Author: saddam213

OnnxStack.StableDiffusion/Common/IStableDiffusionService.cs

@@ -1,4 +1,5 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Config;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
@@ -12,8 +13,36 @@ namespace OnnxStack.StableDiffusion.Common
 {
     public interface IStableDiffusionService
     {
+
+        /// <summary>
+        /// Gets the models.
+        /// </summary>
         List<ModelOptions> Models { get; }
 
+        /// <summary>
+        /// Loads the model.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns></returns>
+        Task<bool> LoadModel(IModelOptions modelOptions);
+
+
+        /// <summary>
+        /// Unloads the model.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns></returns>
+        Task<bool> UnloadModel(IModelOptions modelOptions);
+
+        /// <summary>
+        /// Determines whether the specified model is loaded
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns>
+        ///   <c>true</c> if the specified model is loaded; otherwise, <c>false</c>.
+        /// </returns>
+        bool IsModelLoaded(IModelOptions modelOptions);
+
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
         /// </summary>

OnnxStack.StableDiffusion/Config/ModelOptions.cs

@@ -1,7 +1,5 @@
-using OnnxStack.Common.Config;
-using OnnxStack.StableDiffusion.Common;
+using OnnxStack.StableDiffusion.Common;
 using System.Collections.Immutable;
-using System.Linq;
 
 namespace OnnxStack.StableDiffusion.Config
 {

OnnxStack.StableDiffusion/Services/StableDiffusionService.cs

@@ -24,24 +24,64 @@ public sealed class StableDiffusionService : IStableDiffusionService
         private readonly IDiffuser _imageDiffuser;
         private readonly IDiffuser _inpaintDiffuser;
         private readonly IDiffuser _inpaintLegacyDiffuser;
+        private readonly IOnnxModelService _onnxModelService;
         private readonly StableDiffusionConfig _configuration;
 
+
         /// <summary>
         /// Initializes a new instance of the <see cref="StableDiffusionService"/> class.
         /// </summary>
         /// <param name="schedulerService">The scheduler service.</param>
         public StableDiffusionService(StableDiffusionConfig configuration, IOnnxModelService onnxModelService, IPromptService promptService)
         {
             _configuration = configuration;
+            _onnxModelService = onnxModelService;
             _textDiffuser = new TextDiffuser(onnxModelService, promptService);
             _imageDiffuser = new ImageDiffuser(onnxModelService, promptService);
             _inpaintDiffuser = new InpaintDiffuser(onnxModelService, promptService);
             _inpaintLegacyDiffuser = new InpaintLegacyDiffuser(onnxModelService, promptService);
         }
 
+
+        /// <summary>
+        /// Gets the models.
+        /// </summary>
         public List<ModelOptions> Models => _configuration.OnnxModelSets;
 
 
+        /// <summary>
+        /// Loads the model.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns></returns>
+        public async Task<bool> LoadModel(IModelOptions modelOptions)
+        {
+            var model = await _onnxModelService.LoadModel(modelOptions);
+            return model is not null;
+        }
+
+
+        /// <summary>
+        /// Unloads the model.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns></returns>
+        public async Task<bool> UnloadModel(IModelOptions modelOptions)
+        {
+            return await _onnxModelService.UnloadModel(modelOptions);
+        }
+
+
+        /// <summary>
+        /// Is the model loaded.
+        /// </summary>
+        /// <param name="modelOptions">The model options.</param>
+        /// <returns></returns>
+        public bool IsModelLoaded(IModelOptions modelOptions)
+        {
+            return _onnxModelService.IsModelLoaded(modelOptions);
+        }
+
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
         /// </summary>

OnnxStack.UI/UserControls/ModelPickerControl.xaml.cs

@@ -1,5 +1,5 @@
 using Models;
-using OnnxStack.Core.Services;
+using OnnxStack.StableDiffusion.Common;
 using OnnxStack.UI.Commands;
 using System.ComponentModel;
 using System.Runtime.CompilerServices;
@@ -14,13 +14,13 @@ namespace OnnxStack.UI.UserControls
     /// </summary>
     public partial class ModelPickerControl : UserControl, INotifyPropertyChanged
     {
-        private readonly IOnnxModelService _modelService;
+        private readonly IStableDiffusionService _stableDiffusionService;
 
         /// <summary>Initializes a new instance of the <see cref="ModelPickerControl" /> class.</summary>
         public ModelPickerControl()
         {
             if (!DesignerProperties.GetIsInDesignMode(this))
-                _modelService = App.GetService<IOnnxModelService>();
+                _stableDiffusionService = App.GetService<IStableDiffusionService>();
 
             LoadModelCommand = new AsyncRelayCommand(LoadModel);
             UnloadModelCommand = new AsyncRelayCommand(UnloadModel);
@@ -46,13 +46,13 @@ public ModelOptionsModel SelectedModel
         /// </summary>
         private async Task LoadModel()
         {
-            if (_modelService.IsModelLoaded(SelectedModel.ModelOptions))
+            if (_stableDiffusionService.IsModelLoaded(SelectedModel.ModelOptions))
                 return;
 
             SelectedModel.IsLoading = true;
-            await _modelService.LoadModel(SelectedModel.ModelOptions);
+            var result = await _stableDiffusionService.LoadModel(SelectedModel.ModelOptions);
             SelectedModel.IsLoading = false;
-            SelectedModel.IsLoaded = true;
+            SelectedModel.IsLoaded = result;
         }
 
 
@@ -61,11 +61,11 @@ private async Task LoadModel()
         /// </summary>
         private async Task UnloadModel()
         {
-            if (!_modelService.IsModelLoaded(SelectedModel.ModelOptions))
+            if (!_stableDiffusionService.IsModelLoaded(SelectedModel.ModelOptions))
                 return;
 
             SelectedModel.IsLoading = true;
-            await _modelService.UnloadModel(SelectedModel.ModelOptions);
+            await _stableDiffusionService.UnloadModel(SelectedModel.ModelOptions);
             SelectedModel.IsLoading = false;
             SelectedModel.IsLoaded = false;
         }