Skip to content
This repository was archived by the owner on Nov 27, 2024. It is now read-only.

Commit 6eafa27

Browse files
saddam213saddam213
committed
Realtime Inpaint
1 parent 5d5bf7c commit 6eafa27

File tree

4 files changed

+153
-58
lines changed

4 files changed

+153
-58
lines changed

‎OnnxStack.UI/UserControls/ImageInputControl.xaml‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@
2828
</ContextMenu>
2929
</Image.ContextMenu>
3030
</Image>
31-
<InkCanvas x:Name="MaskCanvas" Background="Transparent" ForceCursor="True" Cursor="Pen" MinHeight="512" MinWidth="512" PreviewMouseLeftButtonDown="MaskCanvas_MouseLeftButtonDown"
31+
<InkCanvas x:Name="MaskCanvas" Background="Transparent" ForceCursor="True" Cursor="Pen" MinHeight="512" MinWidth="512"
32+
PreviewMouseLeftButtonUp="MaskCanvas_MouseLeftButtonUp"
33+
PreviewMouseLeftButtonDown="MaskCanvas_MouseLeftButtonDown"
3234
IsEnabled="{Binding HasResult}"
3335
EditingMode="{Binding MaskEditingMode}"
3436
Width="{Binding SchedulerOptions.Width}"
@@ -56,18 +58,6 @@
5658

5759
<DockPanel >
5860
<StackPanel DockPanel.Dock="Right" Orientation="Horizontal" IsEnabled="{Binding HasResult}" Visibility="{Binding IsMaskEnabled, Converter={StaticResource BooleanToVisibilityConverter}}">
59-
<Button DockPanel.Dock="Right" Command="{Binding SaveMaskCommand}" Width="50" BorderThickness="1,1,0,1">
60-
<userControls:FontAwesome Icon="&#xf0c7;" IconStyle="Light" />
61-
<Button.Style>
62-
<Style TargetType="{x:Type Button}" BasedOn="{StaticResource {x:Type Button}}">
63-
<Style.Triggers>
64-
<DataTrigger Binding="{Binding Path=HasMaskChanged}" Value="True">
65-
<Setter Property="Background" Value="{StaticResource RedHighlight}" />
66-
</DataTrigger>
67-
</Style.Triggers>
68-
</Style>
69-
</Button.Style>
70-
</Button>
7161
<Button DockPanel.Dock="Right" Command="{Binding MaskModeCommand}" Width="50" Visibility="{Binding IsMaskEraserEnabled, Converter={StaticResource BooleanToVisibilityConverter}}">
7262
<userControls:FontAwesome Icon="&#xf304;" IconStyle="Light" />
7363
</Button>

‎OnnxStack.UI/UserControls/ImageInputControl.xaml.cs‎

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using OnnxStack.UI.Services;
66
using System;
77
using System.ComponentModel;
8+
using System.IO;
89
using System.Linq;
910
using System.Runtime.CompilerServices;
1011
using System.Threading.Tasks;
@@ -39,7 +40,6 @@ public ImageInputControl()
3940
LoadImageCommand = new AsyncRelayCommand(LoadImage);
4041
ClearImageCommand = new AsyncRelayCommand(ClearImage);
4142
MaskModeCommand = new AsyncRelayCommand(MaskMode);
42-
SaveMaskCommand = new AsyncRelayCommand(SaveMask);
4343
CopyImageCommand = new AsyncRelayCommand(CopyImage);
4444
PasteImageCommand = new AsyncRelayCommand(PasteImage);
4545
InitializeComponent();
@@ -48,7 +48,6 @@ public ImageInputControl()
4848
public AsyncRelayCommand LoadImageCommand { get; }
4949
public AsyncRelayCommand ClearImageCommand { get; }
5050
public AsyncRelayCommand MaskModeCommand { get; }
51-
public AsyncRelayCommand SaveMaskCommand { get; }
5251
public AsyncRelayCommand CopyImageCommand { get; }
5352
public AsyncRelayCommand PasteImageCommand { get; }
5453
public ImageInput Result
@@ -58,7 +57,11 @@ public ImageInput Result
5857
}
5958

6059
public static readonly DependencyProperty ResultProperty =
61-
DependencyProperty.Register("Result", typeof(ImageInput), typeof(ImageInputControl));
60+
DependencyProperty.Register("Result", typeof(ImageInput), typeof(ImageInputControl), new PropertyMetadata((s, e) =>
61+
{
62+
if (s is ImageInputControl control)
63+
control.SaveMask();
64+
}));
6265

6366
public ImageInput MaskResult
6467
{
@@ -242,6 +245,9 @@ private void UpdateMaskAttributes()
242245
/// <returns></returns>
243246
public BitmapSource CreateMaskImage()
244247
{
248+
if (MaskCanvas.ActualWidth == 0)
249+
return CreateEmptyMaskImage();
250+
245251
// Create a RenderTargetBitmap to render the Canvas content.
246252
var renderBitmap = new RenderTargetBitmap((int)MaskCanvas.ActualWidth, (int)MaskCanvas.ActualHeight, 96, 96, PixelFormats.Pbgra32);
247253

@@ -256,7 +262,27 @@ public BitmapSource CreateMaskImage()
256262
return renderBitmap;
257263
}
258264

259-
private void ShowCropImageDialog(BitmapSource source = null, string sourceFile = null)
265+
266+
public BitmapSource CreateEmptyMaskImage()
267+
{
268+
var wbm = new WriteableBitmap(SchedulerOptions.Width, SchedulerOptions.Height, 96, 96, PixelFormats.Bgra32, null);
269+
BitmapImage bmImage = new BitmapImage();
270+
using (MemoryStream stream = new MemoryStream())
271+
{
272+
PngBitmapEncoder encoder = new PngBitmapEncoder();
273+
encoder.Frames.Add(BitmapFrame.Create(wbm));
274+
encoder.Save(stream);
275+
bmImage.BeginInit();
276+
bmImage.CacheOption = BitmapCacheOption.OnLoad;
277+
bmImage.StreamSource = stream;
278+
bmImage.EndInit();
279+
bmImage.Freeze();
280+
}
281+
return bmImage;
282+
}
283+
284+
285+
private async void ShowCropImageDialog(BitmapSource source = null, string sourceFile = null)
260286
{
261287
try
262288
{
@@ -269,13 +295,14 @@ private void ShowCropImageDialog(BitmapSource source = null, string sourceFile =
269295
loadImageDialog.Initialize(SchedulerOptions.Width, SchedulerOptions.Height, source);
270296
if (loadImageDialog.ShowDialog() == true)
271297
{
272-
ClearImage();
298+
await ClearImage();
273299
Result = new ImageInput
274300
{
275301
Image = loadImageDialog.GetImageResult(),
276302
FileName = loadImageDialog.ImageFile,
277303
};
278304
HasResult = true;
305+
await SaveMask();
279306
}
280307
}
281308

@@ -333,6 +360,16 @@ private void MaskCanvas_MouseLeftButtonDown(object sender, System.Windows.Input.
333360
}
334361

335362

363+
/// <summary>
364+
/// Handles the MouseLeftButtonUp event of the MaskCanvas control.
365+
/// </summary>
366+
/// <param name="sender">The source of the event.</param>
367+
/// <param name="e">The <see cref="MouseButtonEventArgs"/> instance containing the event data.</param>
368+
private async void MaskCanvas_MouseLeftButtonUp(object sender, MouseButtonEventArgs e)
369+
{
370+
await SaveMask();
371+
}
372+
336373
/// <summary>
337374
/// Called on key down.
338375
/// </summary>

‎OnnxStack.UI/Views/ImageInpaint.xaml‎

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,25 +19,24 @@
1919
<Button Content="Cancel" Command="{Binding CancelCommand}" Margin="0,0,1,0"/>
2020
<Button Content="Generate" Command="{Binding GenerateCommand}" IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}" IsDefault="True" Margin="1,0,0,0"/>
2121
</UniformGrid>
22-
<DockPanel IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}">
23-
<UniformGrid Rows="2">
24-
<DockPanel>
25-
<userControls:ModelPickerControl DockPanel.Dock="Top"
26-
UISettings="{Binding UISettings}"
27-
SupportedDiffusers="{Binding SupportedDiffusers}"
28-
SelectedModel="{Binding SelectedModel, Mode=TwoWay}" Models="{Binding RelativeSource={RelativeSource FindAncestor,AncestorType={x:Type Window}}, Path=Models}" />
29-
<userControls:PromptControl
30-
SelectedModel="{Binding SelectedModel}"
31-
PromptOptions="{Binding PromptOptions}"
32-
IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}"/>
33-
</DockPanel>
34-
<userControls:SchedulerControl DockPanel.Dock="Bottom"
22+
<DockPanel>
23+
<userControls:ModelPickerControl DockPanel.Dock="Top"
24+
UISettings="{Binding UISettings}"
25+
SupportedDiffusers="{Binding SupportedDiffusers}"
26+
IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
27+
SelectedModel="{Binding SelectedModel, Mode=TwoWay}" Models="{Binding RelativeSource={RelativeSource FindAncestor,AncestorType={x:Type Window}}, Path=Models}"/>
28+
<UniformGrid Rows="2" IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}">
29+
<userControls:PromptControl
30+
SelectedModel="{Binding SelectedModel}"
31+
PromptOptions="{Binding PromptOptions}"
32+
IsEnabled="{Binding IsControlsEnabled}"/>
33+
<userControls:SchedulerControl DockPanel.Dock="Bottom" Margin="0, 10, 0 ,0"
3534
SelectedModel="{Binding SelectedModel}"
3635
DiffuserType="ImageInpaint"
3736
BatchOptions="{Binding BatchOptions}"
3837
SchedulerOptions="{Binding SchedulerOptions, Mode=TwoWay}"
39-
IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}"
40-
Margin="0, 10, 0 ,0"/>
38+
IsEnabled="{Binding IsControlsEnabled}"
39+
IsAutomationEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"/>
4140
</UniformGrid>
4241
</DockPanel>
4342
</DockPanel>

‎OnnxStack.UI/Views/ImageInpaint.xaml.cs‎

Lines changed: 94 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public partial class ImageInpaint : UserControl, INavigatable, INotifyPropertyCh
3535
private int _selectedTabIndex;
3636
private bool _hasInputResult;
3737
private bool _hasInputMaskResult;
38+
private bool _isControlsEnabled;
3839
private ImageInput _inputImage;
3940
private ImageResult _resultImage;
4041
private ImageInput _inputImageMask;
@@ -65,6 +66,7 @@ public ImageInpaint()
6566
BatchOptions = new BatchOptionsModel();
6667
ImageResults = new ObservableCollection<ImageResult>();
6768
ProgressMax = SchedulerOptions.InferenceSteps;
69+
IsControlsEnabled = true;
6870
InitializeComponent();
6971
}
7072

@@ -167,6 +169,11 @@ public int SelectedTabIndex
167169
set { _selectedTabIndex = value; NotifyPropertyChanged(); }
168170
}
169171

172+
public bool IsControlsEnabled
173+
{
174+
get { return _isControlsEnabled; }
175+
set { _isControlsEnabled = value; NotifyPropertyChanged(); }
176+
}
170177

171178

172179
/// <summary>
@@ -179,9 +186,9 @@ public Task NavigateAsync(ImageResult imageResult)
179186
Reset();
180187
HasResult = false;
181188
ResultImage = null;
189+
InputImageMask = null;
182190
HasInputResult = true;
183191
HasInputMaskResult = false;
184-
InputImageMask = null;
185192
if (imageResult.Model.ModelOptions.Diffusers.Contains(DiffuserType.ImageInpaint)
186193
|| imageResult.Model.ModelOptions.Diffusers.Contains(DiffuserType.ImageInpaintLegacy))
187194
{
@@ -210,24 +217,9 @@ private async Task Generate()
210217
{
211218
HasResult = false;
212219
IsGenerating = true;
220+
IsControlsEnabled = false;
213221
ResultImage = null;
214-
var promptOptions = new PromptOptions
215-
{
216-
Prompt = PromptOptions.Prompt,
217-
NegativePrompt = PromptOptions.NegativePrompt,
218-
DiffuserType = SelectedModel.ModelOptions.Diffusers.Contains(DiffuserType.ImageInpaint)
219-
? DiffuserType.ImageInpaint
220-
: DiffuserType.ImageInpaintLegacy,
221-
InputImage = new StableDiffusion.Models.InputImage
222-
{
223-
ImageBytes = InputImage.Image.GetImageBytes()
224-
},
225-
InputImageMask = new StableDiffusion.Models.InputImage
226-
{
227-
ImageBytes = InputImageMask.Image.GetImageBytes()
228-
}
229-
};
230-
222+
var promptOptions = GetPromptOptions(PromptOptions, InputImage, InputImageMask);
231223
var batchOptions = BatchOptions.ToBatchOptions();
232224
var schedulerOptions = SchedulerOptions.ToSchedulerOptions();
233225
schedulerOptions.Strength = 1; // Make sure strength is 1 for Image Inpainting
@@ -239,8 +231,11 @@ private async Task Generate()
239231
if (resultImage != null)
240232
{
241233
ResultImage = resultImage;
242-
ImageResults.Add(resultImage);
243234
HasResult = true;
235+
if (BatchOptions.IsAutomationEnabled && BatchOptions.DisableHistory)
236+
continue;
237+
238+
ImageResults.Add(resultImage);
244239
}
245240
}
246241
}
@@ -267,8 +262,7 @@ private bool CanExecuteGenerate()
267262
{
268263
return !IsGenerating
269264
&& !string.IsNullOrEmpty(PromptOptions.Prompt)
270-
&& HasInputResult
271-
&& HasInputMaskResult;
265+
&& HasInputResult;
272266
}
273267

274268

@@ -324,6 +318,7 @@ private bool CanExecuteClearHistory()
324318
private void Reset()
325319
{
326320
IsGenerating = false;
321+
IsControlsEnabled = true;
327322
ProgressValue = 0;
328323
}
329324

@@ -338,6 +333,9 @@ private void Reset()
338333
/// <returns></returns>
339334
private async IAsyncEnumerable<ImageResult> ExecuteStableDiffusion(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions)
340335
{
336+
if (!IsExecuteOptionsValid(PromptOptions))
337+
yield break;
338+
341339
_cancelationTokenSource = new CancellationTokenSource();
342340
if (!BatchOptions.IsAutomationEnabled)
343341
{
@@ -347,15 +345,68 @@ private async IAsyncEnumerable<ImageResult> ExecuteStableDiffusion(IModelOptions
347345
}
348346
else
349347
{
350-
var timestamp = Stopwatch.GetTimestamp();
351-
await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
348+
if (_batchOptions.BatchType != BatchOptionType.Realtime)
352349
{
353-
yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
354-
timestamp = Stopwatch.GetTimestamp();
350+
var timestamp = Stopwatch.GetTimestamp();
351+
await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
352+
{
353+
yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
354+
timestamp = Stopwatch.GetTimestamp();
355+
}
356+
yield break;
357+
}
358+
359+
// Realtime Diffusion
360+
IsControlsEnabled = true;
361+
while (!_cancelationTokenSource.IsCancellationRequested)
362+
{
363+
var refreshTimestamp = Stopwatch.GetTimestamp();
364+
if (SchedulerOptions.HasChanged || PromptOptions.HasChanged || HasInputMaskResult)
365+
{
366+
HasInputMaskResult = false;
367+
PromptOptions.HasChanged = false;
368+
SchedulerOptions.HasChanged = false;
369+
var realtimePromptOptions = GetPromptOptions(PromptOptions, InputImage, InputImageMask);
370+
var realtimeSchedulerOptions = SchedulerOptions.ToSchedulerOptions();
371+
372+
var timestamp = Stopwatch.GetTimestamp();
373+
var result = await _stableDiffusionService.GenerateAsBytesAsync(modelOptions, realtimePromptOptions, realtimeSchedulerOptions, RealtimeProgressCallback(), _cancelationTokenSource.Token);
374+
yield return await GenerateResultAsync(result, promptOptions, schedulerOptions, timestamp);
375+
}
376+
await Utils.RefreshDelay(refreshTimestamp, BatchOptions.RealtimeRefreshRate, _cancelationTokenSource.Token);
355377
}
356378
}
357379
}
358380

381+
private bool IsExecuteOptionsValid(PromptOptionsModel prompt)
382+
{
383+
if (string.IsNullOrEmpty(prompt.Prompt))
384+
return false;
385+
386+
return true;
387+
}
388+
389+
390+
private PromptOptions GetPromptOptions(PromptOptionsModel promptOptionsModel, ImageInput imageInput, ImageInput imageInputMask)
391+
{
392+
return new PromptOptions
393+
{
394+
Prompt = promptOptionsModel.Prompt,
395+
NegativePrompt = promptOptionsModel.NegativePrompt,
396+
DiffuserType = SelectedModel.ModelOptions.Diffusers.Contains(DiffuserType.ImageInpaint)
397+
? DiffuserType.ImageInpaint
398+
: DiffuserType.ImageInpaintLegacy,
399+
InputImage = new StableDiffusion.Models.InputImage
400+
{
401+
ImageBytes = imageInput.Image.GetImageBytes()
402+
},
403+
InputImageMask = new StableDiffusion.Models.InputImage
404+
{
405+
ImageBytes = imageInputMask.Image.GetImageBytes()
406+
}
407+
};
408+
}
409+
359410

360411
/// <summary>
361412
/// Generates the result.
@@ -430,6 +481,24 @@ private Action<int, int, int, int> ProgressBatchCallback()
430481
};
431482
}
432483

484+
private Action<int, int> RealtimeProgressCallback()
485+
{
486+
return (value, maximum) =>
487+
{
488+
App.UIInvoke(() =>
489+
{
490+
if (_cancelationTokenSource.IsCancellationRequested)
491+
return;
492+
493+
if (BatchOptions.StepValue != value)
494+
BatchOptions.StepValue = value;
495+
if (BatchOptions.StepsValue != maximum)
496+
BatchOptions.StepsValue = maximum;
497+
});
498+
};
499+
}
500+
501+
433502
#region INotifyPropertyChanged
434503
public event PropertyChangedEventHandler PropertyChanged;
435504
public void NotifyPropertyChanged([CallerMemberName] string property = "")

0 commit comments

Comments
 (0)