ControlNet, ControlNetImage UI support

saddam213 · saddam213 · commit 920554776a09 · 2024-01-05T14:24:10.000+13:00
diff --git a/OnnxStack.StableDiffusion/Config/PromptOptions.cs b/OnnxStack.StableDiffusion/Config/PromptOptions.cs
@@ -20,6 +20,8 @@ public class PromptOptions
 
         public InputImage InputImageMask { get; set; }
 
+        public InputImage InputContolImage { get; set; }
+
         public VideoInput InputVideo { get; set; }
 
         public float VideoInputFPS { get; set; }
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetDiffuser.cs
@@ -193,7 +193,7 @@ protected static DenseTensor<double> CreateConditioningScaleTensor(float conditi
         /// <returns></returns>
         protected DenseTensor<float> PrepareControlImage(PromptOptions promptOptions, SchedulerOptions schedulerOptions)
         {
-            return promptOptions.InputImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);
+            return promptOptions.InputContolImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Services/ImageService.cs b/OnnxStack.StableDiffusion/Services/ImageService.cs
@@ -179,7 +179,7 @@ private async Task<InputImage> GenerateDepthImage(ControlNetModelSet controlNetM
         /// <returns></returns>
         private static InputImage PrepareInputImage(InputImage inputImage, int height, int width)
         {
-            return new InputImage(inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false));
+            return inputImage;
         }
 
 
diff --git a/OnnxStack.UI/Models/SchedulerOptionsModel.cs b/OnnxStack.UI/Models/SchedulerOptionsModel.cs
@@ -33,6 +33,7 @@ public class SchedulerOptionsModel : INotifyPropertyChanged
         private int _originalInferenceSteps = 100;
         private SchedulerType _schedulerType;
         private float _conditioningScale = 0.7f;
+        private bool _processInputImage;
         private bool _hasChanged;
 
         /// <summary>
@@ -208,6 +209,11 @@ public float ConditioningScale
             set { _conditioningScale = value; NotifyPropertyChanged(); }
         }
 
+        public bool ProcessInputImage
+        {
+            get { return _processInputImage; }
+            set { _processInputImage = value; NotifyPropertyChanged(); }
+        }
 
         public bool HasChanged
         {
diff --git a/OnnxStack.UI/Models/StableDiffusionModelSetViewModel.cs b/OnnxStack.UI/Models/StableDiffusionModelSetViewModel.cs
@@ -1,4 +1,5 @@
 ﻿using OnnxStack.StableDiffusion.Config;
+using OnnxStack.StableDiffusion.Enums;
 using System.ComponentModel;
 using System.Runtime.CompilerServices;
 using System.Text.Json.Serialization;
@@ -31,6 +32,8 @@ public bool IsLoading
             set { _isLoading = value; NotifyPropertyChanged(); }
         }
 
+        public bool IsControlNet => ModelSet.Diffusers.Contains(DiffuserType.ControlNet);
+
         public StableDiffusionModelSet ModelSet { get; set; }
 
         #region INotifyPropertyChanged
diff --git a/OnnxStack.UI/UserControls/SchedulerControl.xaml b/OnnxStack.UI/UserControls/SchedulerControl.xaml
@@ -33,7 +33,20 @@
                             <ComboBox ItemsSource="{Binding SchedulerTypes}" SelectedItem="{Binding SchedulerOptions.SchedulerType}"/>
                         </StackPanel>
 
-             
+                        <StackPanel Margin="5,0,0,0" VerticalAlignment="Bottom" Visibility="{Binding SelectedModel.IsControlNet, Converter={StaticResource BooleanToVisibilityConverter}, FallbackValue=Collapsed}">
+                            <CheckBox IsChecked="{Binding SchedulerOptions.ProcessInputImage}"  Content="Process Input Image" Margin="0,0,0,3"/>
+
+                            <!--<StackPanel.Style>
+                                <Style TargetType="{x:Type StackPanel}">
+                                    <Setter Property="IsEnabled" Value="False" />
+                                    <Style.Triggers>
+                                        <DataTrigger Binding="{Binding SchedulerOptions.Strength}" Value="1">
+                                            <Setter Property="IsEnabled" Value="True" />
+                                        </DataTrigger>
+                                    </Style.Triggers>
+                                </Style>
+                            </StackPanel.Style>-->
+                        </StackPanel>
 
 
                     </UniformGrid>
@@ -111,7 +124,7 @@
                                 </i:Interaction.Behaviors>
                             </Slider>
                             <StackPanel.Style>
-                           
+
                                 <Style TargetType="{x:Type StackPanel}">
                                     <Style.Triggers>
                                         <DataTrigger Binding="{Binding Path=DiffuserType}" Value="TextToImage">
@@ -129,7 +142,7 @@
                         </StackPanel>
 
 
-                        <StackPanel Margin="0,0,5,0">
+                        <StackPanel Margin="0,0,5,0" Visibility="{Binding SelectedModel.IsControlNet, Converter={StaticResource BooleanToVisibilityConverter}}">
                             <DockPanel>
                                 <Label>Conditioning Scale</Label>
                                 <TextBlock Text="{Binding ElementName=SliderInitialConditioningScale, Path=Value, StringFormat={}{0:F2}}" VerticalAlignment="Bottom" HorizontalAlignment="Right" FontSize="10" Margin="0,0,6,0" FontWeight="Medium" />
@@ -139,16 +152,7 @@
                                     <behaviors:SliderMouseWheelBehavior />
                                 </i:Interaction.Behaviors>
                             </Slider>
-                            <StackPanel.Style>
-                                <Style TargetType="{x:Type StackPanel}">
-                                    <Setter Property="Visibility" Value="Collapsed" />
-                                    <Style.Triggers>
-                                        <DataTrigger Binding="{Binding Path=DiffuserType}" Value="ControlNet">
-                                            <Setter Property="Visibility" Value="Visible" />
-                                        </DataTrigger>
-                                    </Style.Triggers>
-                                </Style>
-                            </StackPanel.Style>
+
                         </StackPanel>
 
                     </UniformGrid>
diff --git a/OnnxStack.UI/Views/ControlNetView.xaml b/OnnxStack.UI/Views/ControlNetView.xaml
@@ -26,7 +26,7 @@
                          SupportedDiffusers="{Binding SupportedDiffusers}"
                          IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
                          SelectedModel="{Binding SelectedModel, Mode=TwoWay}"/>
-                    <userControls:ControlNetPickerControl 
+                    <userControls:ControlNetPickerControl Visibility="{Binding SelectedModel.IsControlNet, Converter={StaticResource BooleanToVisibilityConverter}}"
                           UISettings="{Binding UISettings}"
                           IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
                           SelectedModel="{Binding SelectedControlNetModel, Mode=TwoWay}"/>
diff --git a/OnnxStack.UI/Views/ImageToImageView.xaml b/OnnxStack.UI/Views/ImageToImageView.xaml
@@ -20,11 +20,19 @@
                 <Button Content="Generate" Command="{Binding GenerateCommand}" IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}" IsDefault="True" Margin="1,0,0,0"/>
             </UniformGrid>
             <DockPanel>
-                <userControls:ModelPickerControl DockPanel.Dock="Top"
-                      UISettings="{Binding UISettings}"
-                      SupportedDiffusers="{Binding SupportedDiffusers}"
-                      IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
-                      SelectedModel="{Binding SelectedModel, Mode=TwoWay}"/>
+                <StackPanel DockPanel.Dock="Top">
+                    <userControls:ModelPickerControl 
+                        UISettings="{Binding UISettings}"
+                        SupportedDiffusers="{Binding SupportedDiffusers}"
+                        IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
+                        SelectedModel="{Binding SelectedModel, Mode=TwoWay}"/>
+                    <userControls:ControlNetPickerControl 
+                          UISettings="{Binding UISettings}"
+                          IsEnabled="{Binding IsGenerating, Converter={StaticResource InverseBoolConverter}}"
+                          SelectedModel="{Binding SelectedControlNetModel, Mode=TwoWay}"
+                          Visibility="{Binding SelectedModel.IsControlNet, Converter={StaticResource BooleanToVisibilityConverter}, FallbackValue=Collapsed}"/>
+                </StackPanel>
+              
                 <DockPanel IsEnabled="{Binding SelectedModel.IsLoaded, FallbackValue=False, TargetNullValue=False}">
                     <userControls:SchedulerControl DockPanel.Dock="Bottom" Margin="0, 10, 0 ,0"
                          SelectedModel="{Binding SelectedModel}"
diff --git a/OnnxStack.UI/Views/ImageToImageView.xaml.cs b/OnnxStack.UI/Views/ImageToImageView.xaml.cs
@@ -26,6 +26,7 @@ public partial class ImageToImageView : UserControl, INavigatable, INotifyProper
     {
         private readonly ILogger<ImageToImageView> _logger;
         private readonly IStableDiffusionService _stableDiffusionService;
+        private readonly IImageService _imageService;
 
         private bool _hasResult;
         private int _progressMax;
@@ -37,6 +38,7 @@ public partial class ImageToImageView : UserControl, INavigatable, INotifyProper
         private ImageInput _inputImage;
         private ImageResult _resultImage;
         private StableDiffusionModelSetViewModel _selectedModel;
+        private ControlNetModelSetViewModel _selectedControlNetModel;
         private PromptOptionsModel _promptOptionsModel;
         private SchedulerOptionsModel _schedulerOptions;
         private CancellationTokenSource _cancelationTokenSource;
@@ -51,9 +53,10 @@ public ImageToImageView()
             {
                 _logger = App.GetService<ILogger<ImageToImageView>>();
                 _stableDiffusionService = App.GetService<IStableDiffusionService>();
+                _imageService = App.GetService<IImageService>();
             }
 
-            SupportedDiffusers = new() { DiffuserType.ImageToImage };
+            SupportedDiffusers = new() { DiffuserType.ImageToImage, DiffuserType.ControlNet };
             CancelCommand = new AsyncRelayCommand(Cancel, CanExecuteCancel);
             GenerateCommand = new AsyncRelayCommand(Generate, CanExecuteGenerate);
             ClearHistoryCommand = new AsyncRelayCommand(ClearHistory, CanExecuteClearHistory);
@@ -84,6 +87,12 @@ public StableDiffusionModelSetViewModel SelectedModel
             get { return _selectedModel; }
             set { _selectedModel = value; NotifyPropertyChanged(); }
         }
+  
+        public ControlNetModelSetViewModel SelectedControlNetModel
+        {
+            get { return _selectedControlNetModel; }
+            set { _selectedControlNetModel = value; NotifyPropertyChanged(); }
+        }
 
         public PromptOptionsModel PromptOptions
         {
@@ -195,13 +204,13 @@ private async Task Generate()
             IsControlsEnabled = false;
             ResultImage = null;
             _cancelationTokenSource = new CancellationTokenSource();
-            var promptOptions = GetPromptOptions(PromptOptions, InputImage);
+            var promptOptions = await GetPromptOptions(PromptOptions, InputImage);
             var schedulerOptions = SchedulerOptions.ToSchedulerOptions();
 
             try
             {
                 var timestamp = Stopwatch.GetTimestamp();
-                var result = await _stableDiffusionService.GenerateAsBytesAsync(new ModelOptions(_selectedModel.ModelSet), promptOptions, schedulerOptions, ProgressCallback(), _cancelationTokenSource.Token);
+                var result = await _stableDiffusionService.GenerateAsBytesAsync(new ModelOptions(_selectedModel.ModelSet, _selectedControlNetModel?.ModelSet), promptOptions, schedulerOptions, ProgressCallback(), _cancelationTokenSource.Token);
                 var resultImage = await GenerateResultAsync(result, promptOptions, schedulerOptions, timestamp);
                 if (resultImage != null)
                 {
@@ -232,8 +241,8 @@ private async Task Generate()
         private bool CanExecuteGenerate()
         {
             return !IsGenerating
-                // && !string.IsNullOrEmpty(PromptOptions.Prompt)
-                && HasInputResult;
+                && HasInputResult
+                && (!SelectedModel.IsControlNet || (SelectedModel.IsControlNet && SelectedControlNetModel?.IsLoaded == true));
         }
 
 
@@ -294,17 +303,38 @@ private void Reset()
         }
 
 
-        private PromptOptions GetPromptOptions(PromptOptionsModel promptOptionsModel, ImageInput imageInput)
+        private async Task<PromptOptions> GetPromptOptions(PromptOptionsModel promptOptionsModel, ImageInput imageInput)
         {
+            if (_selectedModel.IsControlNet)
+            {
+                var controlNetDiffuserType = _schedulerOptions.Strength >= 1
+                     ? DiffuserType.ControlNet
+                     : DiffuserType.ControlNetImage;
+
+                var inputImage = default(InputImage);
+                if (controlNetDiffuserType == DiffuserType.ControlNetImage)
+                    inputImage = new InputImage(imageInput.Image.GetImageBytes());
+
+                var controlImage = new InputImage(imageInput.Image.GetImageBytes());
+                if (_schedulerOptions.ProcessInputImage)
+                    controlImage = await _imageService.PrepareInputImage(_selectedControlNetModel.ModelSet, controlImage, _schedulerOptions.Height, _schedulerOptions.Width);
+
+                return new PromptOptions
+                {
+                    Prompt = promptOptionsModel.Prompt,
+                    NegativePrompt = promptOptionsModel.NegativePrompt,
+                    DiffuserType = controlNetDiffuserType,
+                    InputImage = inputImage,
+                    InputContolImage = controlImage
+                };
+            }
+
             return new PromptOptions
             {
                 Prompt = promptOptionsModel.Prompt,
                 NegativePrompt = promptOptionsModel.NegativePrompt,
                 DiffuserType = DiffuserType.ImageToImage,
-                InputImage = new InputImage
-                {
-                    ImageBytes = imageInput.Image.GetImageBytes()
-                }
+                InputImage = new InputImage(imageInput.Image.GetImageBytes())
             };
         }
 

Original file line number	Diff line number	Diff line change
`@@ -193,7 +193,7 @@ protected static DenseTensor<double> CreateConditioningScaleTensor(float conditi`
`193`	`193`	`/// <returns></returns>`
`194`	`194`	`protected DenseTensor<float> PrepareControlImage(PromptOptions promptOptions, SchedulerOptions schedulerOptions)`
`195`	`195`	`{`
`196`		`- return promptOptions.InputImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);`
	`196`	`+ return promptOptions.InputContolImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);`
`197`	`197`	`}`
`198`	`198`
`199`	`199`
Original file line number	Diff line number	Diff line change
`@@ -179,7 +179,7 @@ private async Task<InputImage> GenerateDepthImage(ControlNetModelSet controlNetM`
`179`	`179`	`/// <returns></returns>`
`180`	`180`	`private static InputImage PrepareInputImage(InputImage inputImage, int height, int width)`
`181`	`181`	`{`
`182`		`- return new InputImage(inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false));`
	`182`	`+ return inputImage;`
`183`	`183`	`}`
`184`	`184`
`185`	`185`