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

Commit 985b2af

Browse files
committed
Realtime TextToImage
1 parent 3753452 commit 985b2af

File tree

6 files changed

+119
-24
lines changed

6 files changed

+119
-24
lines changed

OnnxStack.StableDiffusion/Enums/BatchOptionType.cs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ public enum BatchOptionType
66
Step = 1,
77
Guidance = 2,
88
Strength = 3,
9-
Scheduler = 4
9+
Scheduler = 4,
10+
Realtime = 5
1011
}
1112
}

OnnxStack.UI/Models/BatchOptionsModel.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ public class BatchOptionsModel : INotifyPropertyChanged
1515
private int _stepsValue = 1;
1616
private int _batchValue;
1717
private int _batchsValue = 1;
18+
private int _realtimeRefreshRate = 1000;
1819

1920
public BatchOptionType BatchType
2021
{
@@ -69,6 +70,12 @@ public int BatchsValue
6970
get { return _batchsValue; }
7071
set { _batchsValue = value; NotifyPropertyChanged(); }
7172
}
73+
74+
public int RealtimeRefreshRate
75+
{
76+
get { return _realtimeRefreshRate; }
77+
set { _realtimeRefreshRate = value; NotifyPropertyChanged(); }
78+
}
7279

7380

7481
#region INotifyPropertyChanged

OnnxStack.UI/UserControls/SchedulerControl.xaml

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -239,7 +239,7 @@
239239
<TextBlock Text="Automation" Margin="5,0,0,0"/>
240240
</StackPanel>
241241
</TabItem.Header>
242-
<StackPanel Margin="5">
242+
<StackPanel Margin="5" IsEnabled="{Binding IsAutomationEnabled}">
243243

244244
<CheckBox Content="Enable Automation" IsChecked="{Binding BatchOptions.IsAutomationEnabled}" Margin="0,10,0,0"/>
245245

@@ -349,6 +349,27 @@
349349
</StackPanel.Style>
350350
</StackPanel>
351351

352+
353+
<!--Realtime-->
354+
<StackPanel Margin="0,10">
355+
<UniformGrid Columns="3">
356+
<TextBlock Text="Refresh Rate (ms)"/>
357+
</UniformGrid>
358+
<UniformGrid Columns="2">
359+
<TextBox Text="{Binding BatchOptions.RealtimeRefreshRate}" Margin="0,0,2,0"/>
360+
</UniformGrid>
361+
<StackPanel.Style>
362+
<Style TargetType="{x:Type StackPanel}">
363+
<Setter Property="Visibility" Value="Collapsed" />
364+
<Style.Triggers>
365+
<DataTrigger Binding="{Binding BatchOptions.BatchType, ElementName=UI}" Value="Realtime" >
366+
<Setter Property="Visibility" Value="Visible" />
367+
</DataTrigger>
368+
</Style.Triggers>
369+
</Style>
370+
</StackPanel.Style>
371+
</StackPanel>
372+
352373
<!--Progress-->
353374
<StackPanel Margin="0,5">
354375
<StackPanel Orientation="Horizontal" HorizontalAlignment="Center">

OnnxStack.UI/UserControls/SchedulerControl.xaml.cs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,14 @@ public BatchOptionsModel BatchOptions
9191
DependencyProperty.Register("BatchOptions", typeof(BatchOptionsModel), typeof(SchedulerControl));
9292

9393

94+
public bool IsAutomationEnabled
95+
{
96+
get { return (bool)GetValue(IsAutomationEnabledProperty); }
97+
set { SetValue(IsAutomationEnabledProperty, value); }
98+
}
99+
public static readonly DependencyProperty IsAutomationEnabledProperty =
100+
DependencyProperty.Register("IsAutomationEnabled", typeof(bool), typeof(SchedulerControl));
101+
94102

95103

96104

OnnxStack.UI/Views/TextToImage.xaml

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,24 +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"
35-
SelectedModel="{Binding SelectedModel}"
36-
DiffuserType="TextToImage"
37-
BatchOptions="{Binding BatchOptions}"
38-
SchedulerOptions="{Binding SchedulerOptions, Mode=TwoWay}"
39-
IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}" Margin="0, 10, 0 ,0"/>
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"
34+
SelectedModel="{Binding SelectedModel}"
35+
DiffuserType="TextToImage"
36+
BatchOptions="{Binding BatchOptions}"
37+
SchedulerOptions="{Binding SchedulerOptions, Mode=TwoWay}"
38+
IsEnabled="{Binding IsControlsEnabled}"
39+
IsAutomationEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"/>
4040
</UniformGrid>
4141
</DockPanel>
4242
</DockPanel>

OnnxStack.UI/Views/TextToImage.xaml.cs

Lines changed: 62 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public partial class TextToImageView : UserControl, INavigatable, INotifyPropert
3535
private int _progressValue;
3636
private bool _isGenerating;
3737
private int _selectedTabIndex;
38+
private bool _isControlsEnabled;
3839
private ImageResult _resultImage;
3940
private ModelOptionsModel _selectedModel;
4041
private PromptOptionsModel _promptOptionsModel;
@@ -62,6 +63,7 @@ public TextToImageView()
6263
BatchOptions = new BatchOptionsModel();
6364
ImageResults = new ObservableCollection<ImageResult>();
6465
ProgressMax = SchedulerOptions.InferenceSteps;
66+
IsControlsEnabled = true;
6567
InitializeComponent();
6668
}
6769

@@ -141,6 +143,11 @@ public int SelectedTabIndex
141143
}
142144

143145

146+
public bool IsControlsEnabled
147+
{
148+
get { return _isControlsEnabled; }
149+
set { _isControlsEnabled = value; NotifyPropertyChanged(); }
150+
}
144151

145152

146153
/// <summary>
@@ -175,6 +182,7 @@ private async Task Generate()
175182
{
176183
HasResult = false;
177184
IsGenerating = true;
185+
IsControlsEnabled = false;
178186
ResultImage = null;
179187
var promptOptions = new PromptOptions
180188
{
@@ -274,6 +282,7 @@ private bool CanExecuteClearHistory()
274282
private void Reset()
275283
{
276284
IsGenerating = false;
285+
IsControlsEnabled = true;
277286
ProgressValue = 0;
278287
}
279288

@@ -295,16 +304,65 @@ private async IAsyncEnumerable<ImageResult> ExecuteStableDiffusion(IModelOptions
295304
}
296305
else
297306
{
298-
var timestamp = Stopwatch.GetTimestamp();
299-
await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
307+
if (_batchOptions.BatchType != BatchOptionType.Realtime)
300308
{
301-
yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
302-
timestamp = Stopwatch.GetTimestamp();
309+
var timestamp = Stopwatch.GetTimestamp();
310+
await foreach (var batchResult in _stableDiffusionService.GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, ProgressBatchCallback(), _cancelationTokenSource.Token))
311+
{
312+
yield return await GenerateResultAsync(batchResult.ImageResult.ToImageBytes(), promptOptions, batchResult.SchedulerOptions, timestamp);
313+
timestamp = Stopwatch.GetTimestamp();
314+
}
315+
yield break;
316+
}
317+
318+
// Realtime Diffusion
319+
IsControlsEnabled = true;
320+
while (!_cancelationTokenSource.IsCancellationRequested)
321+
{
322+
var refreshTimestamp = Stopwatch.GetTimestamp();
323+
var realtimePromptOptions = GetPromptOptions(PromptOptions);
324+
var realtimeSchedulerOptions = SchedulerOptions.ToSchedulerOptions();
325+
if (IsLiveOptionsValid(realtimePromptOptions))
326+
{
327+
var timestamp = Stopwatch.GetTimestamp();
328+
var result = await _stableDiffusionService.GenerateAsBytesAsync(modelOptions, realtimePromptOptions, realtimeSchedulerOptions, ProgressCallback(), _cancelationTokenSource.Token);
329+
yield return await GenerateResultAsync(result, promptOptions, schedulerOptions, timestamp);
330+
}
331+
await RealtimeRefreshDelay(refreshTimestamp, _cancelationTokenSource.Token);
303332
}
304333
}
305334
}
306335

307336

337+
private bool IsLiveOptionsValid(PromptOptions prompt)
338+
{
339+
if (string.IsNullOrEmpty(prompt.Prompt))
340+
return false;
341+
342+
return true;
343+
}
344+
345+
346+
private PromptOptions GetPromptOptions(PromptOptionsModel promptOptionsModel)
347+
{
348+
return new PromptOptions
349+
{
350+
Prompt = promptOptionsModel.Prompt,
351+
NegativePrompt = promptOptionsModel.NegativePrompt,
352+
DiffuserType = DiffuserType.TextToImage
353+
};
354+
}
355+
356+
357+
private async Task RealtimeRefreshDelay(long startTime, CancellationToken cancellationToken)
358+
{
359+
var endTime = Stopwatch.GetTimestamp();
360+
var elapsedMilliseconds = (endTime - startTime) * 1000.0 / Stopwatch.Frequency;
361+
int adjustedDelay = Math.Max(0, BatchOptions.RealtimeRefreshRate - (int)elapsedMilliseconds);
362+
await Task.Delay(adjustedDelay, cancellationToken).ConfigureAwait(false);
363+
}
364+
365+
308366
private async Task<ImageResult> GenerateResultAsync(byte[] imageBytes, PromptOptions promptOptions, SchedulerOptions schedulerOptions, long timestamp)
309367
{
310368
var image = Utils.CreateBitmap(imageBytes);

0 commit comments

Comments
 (0)