Merge pull request #100 from saddam213/ControlNet

ControlNet Improvements
TensorStack-AI · Jan 18, 2024 · 90bb9f1 · 90bb9f1
2 parents 3c4e470 + 68fa37f
commit 90bb9f1
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 9 deletions.
diff --git a/OnnxStack.StableDiffusion/Config/ControlNetModelSet.cs b/OnnxStack.StableDiffusion/Config/ControlNetModelSet.cs
@@ -8,6 +8,7 @@ namespace OnnxStack.StableDiffusion.Config
     public record ControlNetModelSet : IOnnxModelSetConfig
     {
         public ControlNetType Type { get; set; }
+        public DiffuserPipelineType PipelineType { get; set; }
         public string Name { get; set; }
         public bool IsEnabled { get; set; }
         public int DeviceId { get; set; }

diff --git a/OnnxStack.UI/Dialogs/AddControlNetModelDialog.xaml b/OnnxStack.UI/Dialogs/AddControlNetModelDialog.xaml
@@ -18,11 +18,17 @@
     <DockPanel DataContext="{Binding ElementName=UI}" Margin="15, 15, 15, 10">
         <StackPanel DockPanel.Dock="Top">
 
-            <StackPanel>
-                <TextBlock Text="Model Type"/>
-                <ComboBox ItemsSource="{Binding Source={StaticResource ControlNetType}}" SelectedItem="{Binding SelectedControlNetType}" />
-            </StackPanel>
+            <UniformGrid Columns="2">
+                <StackPanel>
+                    <TextBlock Text="Model Type"/>
+                    <ComboBox ItemsSource="{Binding Source={StaticResource ControlNetType}}" SelectedItem="{Binding SelectedControlNetType}" />
+                </StackPanel>
 
+                <StackPanel Margin="5,0,0,0">
+                    <TextBlock Text="Pipeline Type"/>
+                    <ComboBox ItemsSource="{Binding Source={StaticResource DiffuserPipelineType}}" SelectedItem="{Binding SelectedPipelineType}" />
+                </StackPanel>
+            </UniformGrid>
 
             <StackPanel Margin="0,10,0,0">
                 <TextBlock Text="Control Model File"/>
@@ -40,7 +46,7 @@
             </StackPanel>
 
             <StackPanel Margin="5">
-             
+
 
                 <ItemsControl ItemsSource="{Binding ValidationResults}">
                     <ItemsControl.ItemTemplate>
@@ -65,7 +71,7 @@
                         </DataTemplate>
                     </ItemsControl.ItemTemplate>
                 </ItemsControl>
-               
+
             </StackPanel>
 
 

diff --git a/OnnxStack.UI/Dialogs/AddControlNetModelDialog.xaml.cs b/OnnxStack.UI/Dialogs/AddControlNetModelDialog.xaml.cs
@@ -83,6 +83,15 @@ public ControlNetType SelectedControlNetType
 
         }
 
+        private DiffuserPipelineType _selectedPipelineType;
+
+        public DiffuserPipelineType SelectedPipelineType
+        {
+            get { return _selectedPipelineType; }
+            set { _selectedPipelineType = value; NotifyPropertyChanged(); CreateModelSet(); }
+        }
+
+
         public ControlNetModelSet ModelSetResult
         {
             get { return _modelSetResult; }
@@ -102,7 +111,7 @@ private void CreateModelSet()
             if (string.IsNullOrEmpty(_modelFile))
                 return;
 
-            _modelSetResult = _modelFactory.CreateControlNetModelSet(ModelName.Trim(), _selectedControlNetType, _modelFile, _annotationModelFile);
+            _modelSetResult = _modelFactory.CreateControlNetModelSet(ModelName.Trim(), _selectedControlNetType, _selectedPipelineType, _modelFile, _annotationModelFile);
 
             // Validate
             ValidationResults.Add(new ValidationResult("Name", !_invalidOptions.Contains(_modelName, StringComparer.OrdinalIgnoreCase) && _modelName.Length > 2 && _modelName.Length < 50));

diff --git a/OnnxStack.UI/Models/UpdateControlNetModelSetViewModel.cs b/OnnxStack.UI/Models/UpdateControlNetModelSetViewModel.cs
@@ -21,6 +21,8 @@ public class UpdateControlNetModelSetViewModel : INotifyPropertyChanged
         private ExecutionProvider _executionProvider;
         private ObservableCollection<ModelFileViewModel> _modelFiles;
         private ControlNetType _controlNetType;
+        private DiffuserPipelineType _pipelineType;
+
 
         public string Name
         {
@@ -33,6 +35,13 @@ public ControlNetType ControlNetType
             get { return _controlNetType; }
             set { _controlNetType = value; NotifyPropertyChanged(); }
         }
+
+        public DiffuserPipelineType PipelineType
+        {
+            get { return _pipelineType; }
+            set { _pipelineType = value; NotifyPropertyChanged(); }
+        }
+
         public int DeviceId
         {
             get { return _deviceId; }
@@ -76,6 +85,7 @@ public static UpdateControlNetModelSetViewModel FromModelSet(ControlNetModelSet
             {
                 Name = modelset.Name,
                 ControlNetType = modelset.Type,
+                PipelineType = modelset.PipelineType,
                 DeviceId = modelset.DeviceId,
                 ExecutionMode = modelset.ExecutionMode,
                 ExecutionProvider = modelset.ExecutionProvider,
@@ -107,6 +117,7 @@ public static ControlNetModelSet ToModelSet(UpdateControlNetModelSetViewModel mo
                 IsEnabled = true,
                 Name = modelset.Name,
                 Type = modelset.ControlNetType,
+                PipelineType = modelset.PipelineType,
                 DeviceId = modelset.DeviceId,
                 ExecutionMode = modelset.ExecutionMode,
                 ExecutionProvider = modelset.ExecutionProvider,

diff --git a/OnnxStack.UI/Services/IModelFactory.cs b/OnnxStack.UI/Services/IModelFactory.cs
@@ -12,6 +12,6 @@ public interface IModelFactory
 
         UpscaleModelSet CreateUpscaleModelSet(string name, string filename, UpscaleModelTemplate modelTemplate);
         StableDiffusionModelSet CreateStableDiffusionModelSet(string name, string folder, StableDiffusionModelTemplate modelTemplate);
-        ControlNetModelSet CreateControlNetModelSet(string name, ControlNetType controlNetType, string modelFilename, string annotationFilename);
+        ControlNetModelSet CreateControlNetModelSet(string name, ControlNetType controlNetType, DiffuserPipelineType pipelineType, string modelFilename, string annotationFilename);
     }
 }
diff --git a/OnnxStack.UI/Services/ModelFactory.cs b/OnnxStack.UI/Services/ModelFactory.cs
@@ -152,7 +152,7 @@ public UpscaleModelSet CreateUpscaleModelSet(string name, string filename, Upsca
         }
 
 
-        public ControlNetModelSet CreateControlNetModelSet(string name, ControlNetType controlNetType, string modelFilename, string annotationFilename)
+        public ControlNetModelSet CreateControlNetModelSet(string name, ControlNetType controlNetType, DiffuserPipelineType pipelineType, string modelFilename, string annotationFilename)
         {
             var models = new List<OnnxModelConfig> { new OnnxModelConfig { Type = OnnxModelType.ControlNet, OnnxModelPath = modelFilename } };
             if (!string.IsNullOrEmpty(annotationFilename))
@@ -162,6 +162,7 @@ public ControlNetModelSet CreateControlNetModelSet(string name, ControlNetType c
             {
                 Name = name,
                 Type = controlNetType,
+                PipelineType = pipelineType,
                 ModelConfigurations = models,
 
                 IsEnabled = true,