From b7e5ee7232ea5e3277e939becbc626f426f138fe Mon Sep 17 00:00:00 2001
From: DarthAffe <darthaffe@wyrez.org>
Date: Fri, 3 Jan 2025 14:56:27 +0100
Subject: [PATCH] Added flash attention parameter to example

---
 Examples/ImageCreationUI/MainWindow.xaml        |  5 +++++
 Examples/ImageCreationUI/MainWindowViewModel.cs | 11 +++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Examples/ImageCreationUI/MainWindow.xaml b/Examples/ImageCreationUI/MainWindow.xaml
index 13aa12f..a549ca1 100644
--- a/Examples/ImageCreationUI/MainWindow.xaml
+++ b/Examples/ImageCreationUI/MainWindow.xaml
@@ -91,6 +91,11 @@
                 <Label Content="Schedule" />
                 <ComboBox ItemsSource="{Binding Source={StaticResource ScheduleDataSource}}" SelectedItem="{Binding Schedule}" />
 
+                <StackPanel Orientation="Horizontal" Margin="0,4,0,0">
+                    <Label Content="Flash Attention" />
+                    <CheckBox VerticalAlignment="Center" IsChecked="{Binding FlashAttention}" />
+                </StackPanel>
+
                 <Button Margin="0,8"  Content="Load Model" Command="{Binding LoadModelCommand}" IsEnabled="{Binding IsReady}" />
 
                 <Separator />
diff --git a/Examples/ImageCreationUI/MainWindowViewModel.cs b/Examples/ImageCreationUI/MainWindowViewModel.cs
index c9922c0..12272fc 100644
--- a/Examples/ImageCreationUI/MainWindowViewModel.cs
+++ b/Examples/ImageCreationUI/MainWindowViewModel.cs
@@ -71,6 +71,13 @@ public class MainWindowViewModel : INotifyPropertyChanged
         set => SetProperty(ref _schedule, value);
     }
 
+    private bool _flashAttention = true;
+    public bool FlashAttention
+    {
+        get => _flashAttention;
+        set => SetProperty(ref _flashAttention, value);
+    }
+
     private string _prompt = string.Empty;
     public string Prompt
     {
@@ -242,14 +249,14 @@ public class MainWindowViewModel : INotifyPropertyChanged
                 restoreDefaultParameters = _model?.ModelParameter.DiffusionModelType != DiffusionModelType.StableDiffusion;
 
                 LogLine($"Loading stable diffusion-model '{ModelPath}'");
-                _model = await Task.Run(() => ModelBuilder.StableDiffusion(ModelPath).WithMultithreading().WithVae(VaePath).WithSchedule(Schedule).Build());
+                _model = await Task.Run(() => ModelBuilder.StableDiffusion(ModelPath).WithMultithreading().WithVae(VaePath).WithSchedule(Schedule).WithFlashAttention(FlashAttention).Build());
             }
             else if (IsFluxSelected)
             {
                 restoreDefaultParameters = _model?.ModelParameter.DiffusionModelType != DiffusionModelType.Flux;
 
                 LogLine($"Loading flux-model '{DiffusionModelPath}'");
-                _model = await Task.Run(() => ModelBuilder.Flux(DiffusionModelPath, ClipLPath, T5xxlPath, VaePath).WithMultithreading().WithSchedule(Schedule).Build());
+                _model = await Task.Run(() => ModelBuilder.Flux(DiffusionModelPath, ClipLPath, T5xxlPath, VaePath).WithMultithreading().WithSchedule(Schedule).WithFlashAttention(FlashAttention).Build());
             }
             else
             {