v0.4.8 - Model weight caching

Author: saddam213

DiffuseApp/DiffuseApp.Common/Message/PipelineRequest.cs

@@ -21,6 +21,12 @@ public PipelineRequest(PipelineConfig config)
             Type = RequestType.PipelineLoad;
         }
 
+        public PipelineRequest(PipelineReloadOptions options)
+        {
+            PipelineReloadOptions = options;
+            Type = RequestType.PipelineReload;
+        }
+
         public PipelineRequest(EnvironmentConfig config, bool isRebuild, bool isReinstall)
         {
             Environment = new EnvironmentRequest
@@ -45,6 +51,7 @@ public PipelineRequest(PipelineOptions options)
         public RequestType Type { get; init; }
         public EnvironmentRequest Environment { get; set; }
         public PipelineConfig PipelineConfig { get; set; }
+        public PipelineReloadOptions PipelineReloadOptions { get; set; }
         public PipelineOptions PipelineOptions { get; set; }
         public int ImageTensorCount { get; set; }
         public int ControlNetTensorCount { get; set; }

DiffuseApp/DiffuseApp.Common/Message/RequestType.cs

@@ -7,6 +7,7 @@ internal enum RequestType
         Environment = 2,
         PipelineLoad = 3,
         PipelineUnload = 4,
-        PipelineRun = 5
+        PipelineRun = 5,
+        PipelineReload = 6
     }
 }

DiffuseApp/DiffuseApp.Common/PipelineClient.cs

@@ -112,6 +112,26 @@ public async Task LoadAsync(PipelineConfig pipeline, CancellationToken cancellat
         }
 
 
+        /// <summary>
+        /// Reload as an asynchronous operation.
+        /// </summary>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        public async Task ReloadAsync(PipelineReloadOptions options, CancellationToken cancellationToken = default)
+        {
+            try
+            {
+                _isCanceled = false;
+                await SendPipelineRequestAsync(new PipelineRequest(options), cancellationToken);
+            }
+            catch (OperationCanceledException)
+            {
+                await KillServerAsync();
+                throw;
+            }
+        }
+
+
         /// <summary>
         /// Unload the PythonPipeline
         /// </summary>

DiffuseApp/DiffuseApp.Common/PipelineServer.cs

@@ -120,6 +120,10 @@ private async Task StartPipelineChannelAsync(CancellationToken cancellationToken
                             {
                                 pipeline = await LoadPipelineAsync(request, cancellationToken);
                             }
+                            else if (request.Type == RequestType.PipelineReload)
+                            {
+                                await ReloadPipelineAsync(request, pipeline, cancellationToken);
+                            }
                             else if (request.Type == RequestType.PipelineUnload)
                             {
                                 await UnloadPipelineAsync(request, pipeline, cancellationToken);
@@ -307,6 +311,32 @@ private async Task<PythonPipeline> LoadPipelineAsync(PipelineRequest request, Ca
         }
 
 
+        /// <summary>
+        /// Reload the pipeline
+        /// </summary>
+        /// <param name="request">The request.</param>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        private async Task<PythonPipeline> ReloadPipelineAsync(PipelineRequest request, PythonPipeline pipeline, CancellationToken cancellationToken)
+        {
+            try
+            {
+                CallbackMessage("Reloading Pipeline...", "Initialize");
+                await pipeline.ReloadAsync(request.PipelineReloadOptions);
+                await _pipelineChannel.SendResponse(cancellationToken);
+
+                CallbackMessage(string.Empty, "Initialize");
+                return pipeline;
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "[PipelineServer] [ReloadPipeline] An exception occurred reloading pipeline.");
+                await _pipelineChannel.SendMessage(new PipelineResponse(ex), cancellationToken);
+                return default;
+            }
+        }
+
+
         /// <summary>
         /// Unloads the pipeline
         /// </summary>

DiffuseApp/DiffuseApp/Common/MemoryProfileModel.cs

@@ -0,0 +1,24 @@
+using Diffuse.Common;
+using TensorStack.WPF;
+
+namespace DiffuseApp.Common
+{
+    public class MemoryProfileModel : BaseModel
+    {
+        private int _memoryGB;
+        private MemoryMode _memoryMode;
+
+        public MemoryMode MemoryMode
+        {
+            get { return _memoryMode; }
+            set { SetProperty(ref _memoryMode, value); }
+        }
+        public int MemoryGB
+        {
+            get { return _memoryGB; }
+            set { SetProperty(ref _memoryGB, value); }
+        }
+    }
+
+
+}

DiffuseApp/DiffuseApp/Common/PipelineModel.cs

@@ -71,16 +71,24 @@ public DataType DataType
         }
 
 
-        public bool IsReloadRequired(PipelineModel pipeline)
+        public bool IsLoadRequired(PipelineModel pipeline)
         {
             return pipeline is null
                 || pipeline.DiffusionModel != _diffusionModel
-                || pipeline.LoraAdapterModel.HasChanged(_loraAdapterModel)
-                || pipeline.ControlNetModel != _controlNetModel
                 || pipeline.MemoryMode != _memoryMode
-                || pipeline.DataType != _dataType
-                || pipeline.ProcessType != _processType;
+                || pipeline.DataType != _dataType;
         }
 
+
+        public bool IsReloadRequired(PipelineModel pipeline)
+        {
+            if (pipeline is null || pipeline.DiffusionModel != _diffusionModel)
+                return false;
+
+            // ProcessType, LoraAdapters and ControlNet are the only options that can be modified
+            return pipeline.ProcessType != _processType
+                || pipeline.ControlNetModel != _controlNetModel
+                || pipeline.LoraAdapterModel.HasChanged(_loraAdapterModel);
+        }
     }
 }

DiffuseApp/DiffuseApp/Controls/DiffusionModelControl.xaml.cs

@@ -1,4 +1,5 @@
 using Diffuse.Common;
+using DiffuseApp.Common;
 using System;
 using System.Collections.ObjectModel;
 using System.Linq;
@@ -566,24 +567,56 @@ public bool HasLoraChanged()
             return _currentLora.HasChanged(LoraAdapters);
         }
 
-    }
-
-    public class MemoryProfileModel : BaseModel
-    {
-        private int _memoryGB;
-        private MemoryMode _memoryMode;
 
-        public MemoryMode MemoryMode
+        public void SetPipeline(PipelineModel pipeline)
         {
-            get { return _memoryMode; }
-            set { SetProperty(ref _memoryMode, value); }
-        }
-        public int MemoryGB
-        {
-            get { return _memoryGB; }
-            set { SetProperty(ref _memoryGB, value); }
+            if (pipeline == null)
+                return;
+
+            if (!ModelCollectionView.Contains(pipeline.DiffusionModel))
+                return;
+
+            SelectedDevice = pipeline.Device;
+            SelectedModel = pipeline.DiffusionModel;
+
+            SelectedDataType = pipeline.DataType;
+            SelectedMemoryMode = SelectedMemoryMode = MemoryModes.FirstOrDefault(x => x.MemoryMode == pipeline.MemoryMode);
+
+            if (IsUpscalerSupported)
+            {
+                IsUpscalerEnabled = pipeline.UpscaleModel is not null;
+                if (pipeline.UpscaleModel is not null)
+                    SelectedUpscaler = pipeline.UpscaleModel;
+            }
+
+            if (IsExtractorSupported)
+            {
+                IsExtractorEnabled = pipeline.ExtractModel is not null;
+                if (pipeline.ExtractModel is not null)
+                    SelectedExtractor = pipeline.ExtractModel;
+            }
+
+            if (IsControlNetSupported)
+            {
+                if (pipeline.ControlNetModel is not null)
+                    SelectedControlNet = pipeline.ControlNetModel;
+            }
+
+            if (IsLoraSupported)
+            {
+                IsLoraEnabled = !pipeline.LoraAdapterModel.IsNullOrEmpty();
+                if (IsLoraEnabled)
+                {
+                    LoraAdapters.Clear();
+                    foreach (var loraAdapter in pipeline.LoraAdapterModel)
+                    {
+                        LoraAdapters.Add(loraAdapter);
+                    }
+                }
+            }
+
+            ValidateSelection();
         }
     }
 
-
 }

DiffuseApp/DiffuseApp/Controls/ModelControl.xaml.cs

@@ -161,8 +161,8 @@ private bool CanUnload()
         private bool HasCurrentChanged()
         {
             return _currentDevice != SelectedDevice
-                || _currentExtractor != SelectedExtractor
-                || _currentUpscaler != SelectedUpscaler;
+                || (IsExtractorEnabled && _currentExtractor != SelectedExtractor)
+                || (IsUpscalerEnabled && _currentUpscaler != SelectedUpscaler);
         }
 
 
@@ -240,12 +240,27 @@ private Task OnIsPipelineLoadedChanged()
 
         private void ValidateSelection()
         {
-
             var isExtractValid = !IsExtractorEnabled || ExtractCollectionView?.IsEmpty == false;
             var isUpscaleValid = !IsUpscalerEnabled || UpscaleCollectionView?.IsEmpty == false;
             var isCurrentValid = !HasCurrentChanged();
             IsSelectionValid = isCurrentValid && isExtractValid && isUpscaleValid && IsPipelineLoaded;
             LoadCommand.RaiseCanExecuteChanged();
         }
+
+
+        public void SetPipeline(PipelineModel pipeline)
+        {
+            if (pipeline == null)
+                return;
+
+            SelectedDevice = pipeline.Device;
+            if (IsUpscalerEnabled && pipeline.UpscaleModel is not null)
+                SelectedUpscaler = pipeline.UpscaleModel;
+
+            if (IsExtractorEnabled && pipeline.ExtractModel is not null)
+                SelectedExtractor = pipeline.ExtractModel;
+
+            ValidateSelection();
+        }
     }
 }