Move model templates to ModelFactory

Author: saddam213

OnnxStack.UI/Dialogs/AddModelDialog.xaml

@@ -20,21 +20,8 @@
 
             <StackPanel >
                 <TextBlock Text="Model Type"/>
-                <ComboBox ItemsSource="{Binding ModelTemplates}" SelectedItem="{Binding ModelTemplate}" DisplayMemberPath="PipelineType" >
-                    <ComboBox.ItemContainerStyle>
-                        <Style TargetType="{x:Type ComboBoxItem}" BasedOn="{StaticResource {x:Type ComboBoxItem}}">
-                            <Setter Property="Visibility" Value="Visible" />
-                            <Style.Triggers>
-                                <DataTrigger Binding="{Binding Category}" Value="Upscaler">
-                                    <Setter Property="Visibility" Value="Collapsed"/>
-                                </DataTrigger>
-                                <DataTrigger Binding="{Binding IsUserTemplate}" Value="True">
-                                    <Setter Property="Visibility" Value="Collapsed"/>
-                                </DataTrigger>
-                            </Style.Triggers>
-
-                        </Style>
-                    </ComboBox.ItemContainerStyle>
+                <ComboBox ItemsSource="{Binding ModelTemplates}" SelectedItem="{Binding ModelTemplate}" DisplayMemberPath="Name" >
+                   
                 </ComboBox>
             </StackPanel>
 

OnnxStack.UI/Dialogs/AddModelDialog.xaml.cs

@@ -37,17 +37,11 @@ public AddModelDialog(OnnxStackUIConfig settings, IModelFactory modelFactory, IL
             _modelFactory = modelFactory;
             SaveCommand = new AsyncRelayCommand(Save, CanExecuteSave);
             CancelCommand = new AsyncRelayCommand(Cancel);
-            ModelTemplates = GetModelTemplates().ToList();
+            ModelTemplates = new List<StableDiffusionModelTemplate>( _modelFactory.GetStableDiffusionModelTemplates());
             InvalidOptions = _settings.StableDiffusionModelSets.Select(x => x.Name.ToLower()).ToList();
             InitializeComponent();
         }
 
-        private IEnumerable<StableDiffusionModelTemplate> GetModelTemplates()
-        {
-            yield return new StableDiffusionModelTemplate("StableDiffusion", DiffuserPipelineType.StableDiffusion, ModelType.Base, 512, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
-            yield return new StableDiffusionModelTemplate("SDXL", DiffuserPipelineType.StableDiffusionXL, ModelType.Base, 1024, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
-        }
-
         public AsyncRelayCommand SaveCommand { get; }
         public AsyncRelayCommand CancelCommand { get; }
         public ObservableCollection<ValidationResult> ValidationResults { get; set; } = new ObservableCollection<ValidationResult>();
@@ -105,6 +99,8 @@ private void CreateModelSet()
         {
             _modelSetResult = null;
             ValidationResults.Clear();
+            if (_modelTemplate is null)
+                return;
             if (string.IsNullOrEmpty(_modelFolder))
                 return;
 

OnnxStack.UI/Dialogs/AddUpscaleModelDialog.xaml.cs

@@ -36,16 +36,11 @@ public AddUpscaleModelDialog(OnnxStackUIConfig settings, IModelFactory modelFact
             _modelFactory = modelFactory;
             SaveCommand = new AsyncRelayCommand(Save, CanExecuteSave);
             CancelCommand = new AsyncRelayCommand(Cancel);
-            ModelTemplates = GetModelTemplates().ToList();
+            ModelTemplates = new List<UpscaleModelTemplate>(_modelFactory.GetUpscaleModelTemplates());
             InvalidOptions = _settings.UpscaleModelSets.Select(x => x.Name.ToLower()).ToList();
             InitializeComponent();
         }
 
-        private IEnumerable<UpscaleModelTemplate> GetModelTemplates()
-        {
-            yield return new UpscaleModelTemplate("Upscale 2x", 2, 512);
-            yield return new UpscaleModelTemplate("Upscale 4x", 4, 512);
-        }
 
         public AsyncRelayCommand SaveCommand { get; }
         public AsyncRelayCommand CancelCommand { get; }

OnnxStack.UI/Services/IModelFactory.cs

@@ -1,9 +1,13 @@
 using OnnxStack.StableDiffusion.Config;
+using System.Collections.Generic;
 
 namespace OnnxStack.UI.Services
 {
     public interface IModelFactory
     {
+        IEnumerable<UpscaleModelTemplate> GetUpscaleModelTemplates();
+        IEnumerable<StableDiffusionModelTemplate> GetStableDiffusionModelTemplates();
+
         UpscaleModelSet CreateUpscaleModelSet(string name, string filename, UpscaleModelTemplate modelTemplate);
         StableDiffusionModelSet CreateStableDiffusionModelSet(string name, string folder, StableDiffusionModelTemplate modelTemplate);
     }

OnnxStack.UI/Services/ModelFactory.cs

@@ -22,6 +22,30 @@ public ModelFactory(OnnxStackUIConfig settings)
                 _defaultTokenizerPath = defaultTokenizerPath;
         }
 
+        public IEnumerable<UpscaleModelTemplate> GetUpscaleModelTemplates()
+        {
+            yield return new UpscaleModelTemplate("Upscale 2x", 2, 512);
+            yield return new UpscaleModelTemplate("Upscale 4x", 4, 512);
+        }
+
+
+        public IEnumerable<StableDiffusionModelTemplate> GetStableDiffusionModelTemplates()
+        {
+            yield return new StableDiffusionModelTemplate("SD", DiffuserPipelineType.StableDiffusion, ModelType.Base, 512, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+            yield return new StableDiffusionModelTemplate("SD-Inpaint", DiffuserPipelineType.StableDiffusion, ModelType.Base, 512,  DiffuserType.ImageInpaint);
+
+            yield return new StableDiffusionModelTemplate("SDXL", DiffuserPipelineType.StableDiffusionXL, ModelType.Base, 1024, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+            yield return new StableDiffusionModelTemplate("SDXL-Inpaint", DiffuserPipelineType.StableDiffusionXL, ModelType.Base, 1024, DiffuserType.ImageInpaint);
+            yield return new StableDiffusionModelTemplate("SDXL-Refiner", DiffuserPipelineType.StableDiffusionXL, ModelType.Refiner, 1024, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+            yield return new StableDiffusionModelTemplate("SDXL-Turbo", DiffuserPipelineType.StableDiffusionXL, ModelType.Base, 512, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+
+            yield return new StableDiffusionModelTemplate("LCM", DiffuserPipelineType.LatentConsistency, ModelType.Base, 512, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+            yield return new StableDiffusionModelTemplate("LCM-SDXL", DiffuserPipelineType.LatentConsistencyXL, ModelType.Base, 1024, DiffuserType.TextToImage, DiffuserType.ImageToImage, DiffuserType.ImageInpaintLegacy);
+
+            yield return new StableDiffusionModelTemplate("InstaFlow", DiffuserPipelineType.InstaFlow, ModelType.Base, 512, DiffuserType.TextToImage);
+        }
+
+
         public StableDiffusionModelSet CreateStableDiffusionModelSet(string name, string folder, StableDiffusionModelTemplate modelTemplate)
         {
             var modelSet = new StableDiffusionModelSet