11using Microsoft . Extensions . Logging ;
22using Microsoft . Win32 ;
33using Models ;
4+ using OnnxStack . ImageUpscaler . Config ;
45using OnnxStack . StableDiffusion . Config ;
5- using OnnxStack . StableDiffusion . Enums ;
66using OnnxStack . UI . Commands ;
77using OnnxStack . UI . Models ;
88using OnnxStack . UI . Views ;
@@ -27,36 +27,43 @@ public partial class MainWindow : Window, INotifyPropertyChanged
2727 private INavigatable _selectedTabItem ;
2828 private readonly ILogger < MainWindow > _logger ;
2929 private ObservableCollection < ModelOptionsModel > _models ;
30+ private ObservableCollection < UpscaleModelSetModel > _upscaleModels ;
3031
31- public MainWindow ( StableDiffusionConfig configuration , OnnxStackUIConfig uiSettings , ILogger < MainWindow > logger )
32+ public MainWindow ( OnnxStackUIConfig uiSettings , StableDiffusionConfig configuration , ImageUpscalerConfig upscaleConfiguration , ILogger < MainWindow > logger )
3233 {
3334 _logger = logger ;
3435 UISettings = uiSettings ;
35- SaveImageCommand = new AsyncRelayCommand < ImageResult > ( SaveImageFile ) ;
36+ SaveImageCommand = new AsyncRelayCommand < UpscaleResult > ( SaveUpscaleImageFile ) ;
37+ SaveImageResultCommand = new AsyncRelayCommand < ImageResult > ( SaveImageResultFile ) ;
3638 SaveBlueprintCommand = new AsyncRelayCommand < ImageResult > ( SaveBlueprintFile ) ;
3739 NavigateTextToImageCommand = new AsyncRelayCommand < ImageResult > ( NavigateTextToImage ) ;
3840 NavigateImageToImageCommand = new AsyncRelayCommand < ImageResult > ( NavigateImageToImage ) ;
3941 NavigateImageInpaintCommand = new AsyncRelayCommand < ImageResult > ( NavigateImageInpaint ) ;
4042 NavigateImagePaintToImageCommand = new AsyncRelayCommand < ImageResult > ( NavigateImagePaintToImage ) ;
41-
43+ NavigateUpscalerCommand = new AsyncRelayCommand < ImageResult > ( NavigateUpscaler ) ;
4244 WindowCloseCommand = new AsyncRelayCommand ( WindowClose ) ;
4345 WindowRestoreCommand = new AsyncRelayCommand ( WindowRestore ) ;
4446 WindowMinimizeCommand = new AsyncRelayCommand ( WindowMinimize ) ;
4547 WindowMaximizeCommand = new AsyncRelayCommand ( WindowMaximize ) ;
4648 Models = CreateModelOptions ( configuration . ModelSets ) ;
49+ UpscaleModels = CreateUpscaleModelOptions ( upscaleConfiguration . ModelSets ) ;
4750 InitializeComponent ( ) ;
4851 }
4952
53+
54+
5055 public AsyncRelayCommand WindowMinimizeCommand { get ; }
5156 public AsyncRelayCommand WindowRestoreCommand { get ; }
5257 public AsyncRelayCommand WindowMaximizeCommand { get ; }
5358 public AsyncRelayCommand WindowCloseCommand { get ; }
54- public AsyncRelayCommand < ImageResult > SaveImageCommand { get ; }
59+ public AsyncRelayCommand < UpscaleResult > SaveImageCommand { get ; }
60+ public AsyncRelayCommand < ImageResult > SaveImageResultCommand { get ; }
5561 public AsyncRelayCommand < ImageResult > SaveBlueprintCommand { get ; }
5662 public AsyncRelayCommand < ImageResult > NavigateTextToImageCommand { get ; }
5763 public AsyncRelayCommand < ImageResult > NavigateImageToImageCommand { get ; }
5864 public AsyncRelayCommand < ImageResult > NavigateImageInpaintCommand { get ; }
5965 public AsyncRelayCommand < ImageResult > NavigateImagePaintToImageCommand { get ; }
66+ public AsyncRelayCommand < ImageResult > NavigateUpscalerCommand { get ; }
6067
6168 public OnnxStackUIConfig UISettings
6269 {
@@ -72,6 +79,13 @@ public ObservableCollection<ModelOptionsModel> Models
7279 set { _models = value ; NotifyPropertyChanged ( ) ; }
7380 }
7481
82+ public ObservableCollection < UpscaleModelSetModel > UpscaleModels
83+ {
84+ get { return _upscaleModels ; }
85+ set { _upscaleModels = value ; NotifyPropertyChanged ( ) ; }
86+ }
87+
88+
7589 public int SelectedTabIndex
7690 {
7791 get { return _selectedTabIndex ; }
@@ -104,6 +118,10 @@ private async Task NavigateImagePaintToImage(ImageResult result)
104118 await NavigateToTab ( TabId . PaintToImage , result ) ;
105119 }
106120
121+ private async Task NavigateUpscaler ( ImageResult result )
122+ {
123+ await NavigateToTab ( TabId . Upscaler , result ) ;
124+ }
107125
108126 private async Task NavigateToTab ( TabId tab , ImageResult imageResult )
109127 {
@@ -116,7 +134,8 @@ private enum TabId
116134 TextToImage = 0 ,
117135 ImageToImage = 1 ,
118136 ImageInpaint = 2 ,
119- PaintToImage = 3
137+ PaintToImage = 3 ,
138+ Upscaler = 4
120139 }
121140
122141 private ObservableCollection < ModelOptionsModel > CreateModelOptions ( List < StableDiffusionModelSet > onnxModelSets )
@@ -131,29 +150,47 @@ private ObservableCollection<ModelOptionsModel> CreateModelOptions(List<StableDi
131150 return new ObservableCollection < ModelOptionsModel > ( models ) ;
132151 }
133152
134- private async Task SaveImageFile ( ImageResult imageResult )
153+
154+ private ObservableCollection < UpscaleModelSetModel > CreateUpscaleModelOptions ( List < UpscaleModelSet > modelSets )
155+ {
156+ var models = modelSets . Select ( model => new UpscaleModelSetModel
157+ {
158+ Name = model . Name ,
159+ ModelOptions = model ,
160+ IsEnabled = model . IsEnabled
161+ } ) ;
162+ return new ObservableCollection < UpscaleModelSetModel > ( models ) ;
163+ }
164+
165+ private async Task SaveImageResultFile ( ImageResult imageResult )
135166 {
136167 try
137168 {
138- var saveFileDialog = new SaveFileDialog
139- {
140- Title = Title ,
141- Filter = "png files (*.png)|*.png" ,
142- DefaultExt = "png" ,
143- AddExtension = true ,
144- RestoreDirectory = true ,
145- InitialDirectory = Environment . GetFolderPath ( Environment . SpecialFolder . MyPictures ) ,
146- FileName = $ "image-{ imageResult . SchedulerOptions . Seed } .png"
147- } ;
169+ var filename = GetSaveFilename ( $ "image-{ imageResult . SchedulerOptions . Seed } ") ;
170+ if ( string . IsNullOrEmpty ( filename ) )
171+ return ;
148172
149- var dialogResult = saveFileDialog . ShowDialog ( ) ;
150- if ( dialogResult == false )
151- {
152- _logger . LogInformation ( "Saving image canceled" ) ;
173+ var result = await imageResult . Image . SaveImageFileAsync ( filename ) ;
174+ if ( ! result )
175+ _logger . LogError ( "Error saving image" ) ;
176+
177+ }
178+ catch ( Exception ex )
179+ {
180+ _logger . LogError ( ex , "Error saving image" ) ;
181+ }
182+ }
183+
184+
185+ private async Task SaveUpscaleImageFile ( UpscaleResult imageResult )
186+ {
187+ try
188+ {
189+ var filename = GetSaveFilename ( $ "image-{ imageResult . Info . OutputWidth } x{ imageResult . Info . OutputHeight } ") ;
190+ if ( string . IsNullOrEmpty ( filename ) )
153191 return ;
154- }
155192
156- var result = await imageResult . SaveImageFileAsync ( saveFileDialog . FileName ) ;
193+ var result = await imageResult . Image . SaveImageFileAsync ( filename ) ;
157194 if ( ! result )
158195 _logger . LogError ( "Error saving image" ) ;
159196
@@ -187,6 +224,8 @@ private async Task SaveBlueprintFile(ImageResult imageResult)
187224 return ;
188225 }
189226
227+
228+
190229 var result = await imageResult . SaveBlueprintFileAsync ( saveFileDialog . FileName ) ;
191230 if ( ! result )
192231 _logger . LogError ( "Error saving image blueprint" ) ;
@@ -199,6 +238,30 @@ private async Task SaveBlueprintFile(ImageResult imageResult)
199238 }
200239
201240
241+ private string GetSaveFilename ( string initialFilename )
242+ {
243+ var saveFileDialog = new SaveFileDialog
244+ {
245+ Title = Title ,
246+ Filter = "png files (*.png)|*.png" ,
247+ DefaultExt = "png" ,
248+ AddExtension = true ,
249+ RestoreDirectory = true ,
250+ InitialDirectory = Environment . GetFolderPath ( Environment . SpecialFolder . MyPictures ) ,
251+ FileName = $ "{ initialFilename } .png"
252+ } ;
253+
254+ var dialogResult = saveFileDialog . ShowDialog ( ) ;
255+ if ( dialogResult == false )
256+ {
257+ _logger . LogInformation ( "Saving image canceled" ) ;
258+ return null ;
259+ }
260+
261+ return saveFileDialog . FileName ;
262+ }
263+
264+
202265 /// <summary>
203266 /// Gets or sets the output log.
204267 /// </summary>
0 commit comments