Realtime ImageToImage

Author: saddam213

OnnxStack.UI/Views/ImageToImage.xaml

@@ -19,25 +19,24 @@
                 <Button Content="Cancel" Command="{Binding CancelCommand}" Margin="0,0,1,0"/>
                 <Button Content="Generate" Command="{Binding GenerateCommand}" IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}" IsDefault="True" Margin="1,0,0,0"/>
             </UniformGrid>
-            <DockPanel IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}">
-                <UniformGrid Rows="2">
-                    <DockPanel>
-                        <userControls:ModelPickerControl DockPanel.Dock="Top"
-                             UISettings="{Binding UISettings}"
-                             SupportedDiffusers="{Binding SupportedDiffusers}"
-                             SelectedModel="{Binding SelectedModel, Mode=TwoWay}" Models="{Binding RelativeSource={RelativeSource FindAncestor,AncestorType={x:Type Window}}, Path=Models}"/>
-                        <userControls:PromptControl 
-                             SelectedModel="{Binding SelectedModel}"
-                             PromptOptions="{Binding PromptOptions}" 
-                             IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}"/>
-                    </DockPanel>
-                    <userControls:SchedulerControl  DockPanel.Dock="Bottom"
-                          SelectedModel="{Binding SelectedModel}"
-                          DiffuserType="ImageToImage" 
-                          BatchOptions="{Binding BatchOptions}"
-                          SchedulerOptions="{Binding SchedulerOptions, Mode=TwoWay}" 
-                          IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}" 
-                          Margin="0, 10, 0 ,0"/>
+            <DockPanel>
+                <userControls:ModelPickerControl DockPanel.Dock="Top"
+                      UISettings="{Binding UISettings}"
+                      SupportedDiffusers="{Binding SupportedDiffusers}"
+                      IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
+                      SelectedModel="{Binding SelectedModel, Mode=TwoWay}" Models="{Binding RelativeSource={RelativeSource FindAncestor,AncestorType={x:Type Window}}, Path=Models}"/>
+                <UniformGrid Rows="2" IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}">
+                    <userControls:PromptControl 
+                        SelectedModel="{Binding SelectedModel}"
+                        PromptOptions="{Binding PromptOptions}" 
+                        IsEnabled="{Binding IsControlsEnabled}"/>
+                    <userControls:SchedulerControl DockPanel.Dock="Bottom" Margin="0, 10, 0 ,0"
+                        SelectedModel="{Binding SelectedModel}"
+                        DiffuserType="ImageToImage" 
+                        BatchOptions="{Binding BatchOptions}"
+                        SchedulerOptions="{Binding SchedulerOptions, Mode=TwoWay}" 
+                        IsEnabled="{Binding IsControlsEnabled}"
+                        IsAutomationEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"/>
                 </UniformGrid>
             </DockPanel>
         </DockPanel>

OnnxStack.UI/Views/ImageToImage.xaml.cs

@@ -35,6 +35,7 @@ public partial class ImageToImage : UserControl, INavigatable, INotifyPropertyCh
         private bool _isGenerating;
         private int _selectedTabIndex;
         private bool _hasInputResult;
+        private bool _isControlsEnabled;
         private ImageInput _inputImage;
         private ImageResult _resultImage;
         private ModelOptionsModel _selectedModel;
@@ -64,6 +65,7 @@ public ImageToImage()
             BatchOptions = new BatchOptionsModel();
             ImageResults = new ObservableCollection<ImageResult>();
             ProgressMax = SchedulerOptions.InferenceSteps;
+            IsControlsEnabled = true;
             InitializeComponent();
         }
 
@@ -153,6 +155,12 @@ public int SelectedTabIndex
             set { _selectedTabIndex = value; NotifyPropertyChanged(); }
         }
 
+        public bool IsControlsEnabled
+        {
+            get { return _isControlsEnabled; }
+            set { _isControlsEnabled = value; NotifyPropertyChanged(); }
+        }
+
 
         /// <summary>
         /// Called on Navigate
@@ -192,18 +200,9 @@ private async Task Generate()
         {
             HasResult = false;
             IsGenerating = true;
+            IsControlsEnabled = false;
             ResultImage = null;
-            var promptOptions = new PromptOptions
-            {
-                Prompt = PromptOptions.Prompt,
-                NegativePrompt = PromptOptions.NegativePrompt,
-                DiffuserType = DiffuserType.ImageToImage,
-                InputImage = new StableDiffusion.Models.InputImage
-                {
-                    ImageBytes = InputImage.Image.GetImageBytes()
-                }
-            };
-
+            var promptOptions = GetPromptOptions(PromptOptions, InputImage);
             var batchOptions = BatchOptions.ToBatchOptions();
             var schedulerOptions = SchedulerOptions.ToSchedulerOptions();
 
@@ -214,8 +213,11 @@ private async Task Generate()
                     if (resultImage != null)
                     {
                         ResultImage = resultImage;
-                        ImageResults.Add(resultImage);
                         HasResult = true;
+                        if (BatchOptions.IsAutomationEnabled && BatchOptions.DisableHistory)
+                            continue;
+
+                        ImageResults.Add(resultImage);
                     }
                 }
             }
@@ -241,7 +243,7 @@ private async Task Generate()
         private bool CanExecuteGenerate()
         {
             return !IsGenerating
-                && !string.IsNullOrEmpty(PromptOptions.Prompt) 
+                && !string.IsNullOrEmpty(PromptOptions.Prompt)
                 && HasInputResult;
         }
 
@@ -298,6 +300,7 @@ private bool CanExecuteClearHistory()
         private void Reset()
         {
             IsGenerating = false;
+            IsControlsEnabled = true;
             ProgressValue = 0;
         }
 
@@ -312,6 +315,9 @@ private void Reset()
         /// <returns></returns>
         private async IAsyncEnumerable<ImageResult> ExecuteStableDiffusion(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions)
         {
+            if (!IsExecuteOptionsValid(PromptOptions))
+                yield break;
+
             _cancelationTokenSource = new CancellationTokenSource();
             if (!BatchOptions.IsAutomationEnabled)
             {
@@ -321,15 +327,61 @@ private async IAsyncEnumerable<ImageResult> ExecuteStableDiffusion(IModelOptions
             }
             else
             {
-                var timestamp = Stopwatch.GetTimestamp();
-                await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
+                if (_batchOptions.BatchType != BatchOptionType.Realtime)
+                {
+                    var timestamp = Stopwatch.GetTimestamp();
+                    await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
+                    {
+                        yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
+                        timestamp = Stopwatch.GetTimestamp();
+                    }
+                    yield break;
+                }
+
+                // Realtime Diffusion
+                IsControlsEnabled = true;
+                while (!_cancelationTokenSource.IsCancellationRequested)
                 {
-                    yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
-                    timestamp = Stopwatch.GetTimestamp();
+                    var refreshTimestamp = Stopwatch.GetTimestamp();
+                    if (SchedulerOptions.HasChanged || PromptOptions.HasChanged)
+                    {
+                        PromptOptions.HasChanged = false;
+                        SchedulerOptions.HasChanged = false;
+                        var realtimePromptOptions = GetPromptOptions(PromptOptions, InputImage);
+                        var realtimeSchedulerOptions = SchedulerOptions.ToSchedulerOptions();
+
+                        var timestamp = Stopwatch.GetTimestamp();
+                        var result = await _stableDiffusionService.GenerateAsBytesAsync(modelOptions, realtimePromptOptions, realtimeSchedulerOptions, RealtimeProgressCallback(), _cancelationTokenSource.Token);
+                        yield return await GenerateResultAsync(result, promptOptions, schedulerOptions, timestamp);
+                    }
+                    await Utils.RefreshDelay(refreshTimestamp, BatchOptions.RealtimeRefreshRate, _cancelationTokenSource.Token);
                 }
             }
         }
 
+        private bool IsExecuteOptionsValid(PromptOptionsModel prompt)
+        {
+            if (string.IsNullOrEmpty(prompt.Prompt))
+                return false;
+
+            return true;
+        }
+
+
+        private PromptOptions GetPromptOptions(PromptOptionsModel promptOptionsModel, ImageInput imageInput)
+        {
+            return new PromptOptions
+            {
+                Prompt = promptOptionsModel.Prompt,
+                NegativePrompt = promptOptionsModel.NegativePrompt,
+                DiffuserType = DiffuserType.ImageToImage,
+                InputImage = new StableDiffusion.Models.InputImage
+                {
+                    ImageBytes = imageInput.Image.GetImageBytes()
+                }
+            };
+        }
+
 
         /// <summary>
         /// Generates the result.
@@ -404,6 +456,24 @@ private Action<int, int, int, int> ProgressBatchCallback()
             };
         }
 
+        private Action<int, int> RealtimeProgressCallback()
+        {
+            return (value, maximum) =>
+            {
+                App.UIInvoke(() =>
+                {
+                    if (_cancelationTokenSource.IsCancellationRequested)
+                        return;
+
+                    if (BatchOptions.StepValue != value)
+                        BatchOptions.StepValue = value;
+                    if (BatchOptions.StepsValue != maximum)
+                        BatchOptions.StepsValue = maximum;
+                });
+            };
+        }
+
+
         #region INotifyPropertyChanged
         public event PropertyChangedEventHandler PropertyChanged;
         public void NotifyPropertyChanged([CallerMemberName] string property = "")