SDXL Refiner model support

saddam213 · saddam213 · commit 091cb0dfbf9a · 2023-12-03T18:10:16.000+13:00
diff --git a/OnnxStack.StableDiffusion/Common/IModelOptions.cs b/OnnxStack.StableDiffusion/Common/IModelOptions.cs
@@ -15,6 +15,7 @@ public interface IModelOptions : IOnnxModel
         int TokenizerLimit { get; set; }
         int TokenizerLength { get; set; }
         int Tokenizer2Length { get; set; }
+        ModelType ModelType { get; set; }
         TokenizerType TokenizerType { get; set; }
         DiffuserPipelineType PipelineType { get; set; }
         List<DiffuserType> Diffusers { get; set; }
diff --git a/OnnxStack.StableDiffusion/Config/ModelOptions.cs b/OnnxStack.StableDiffusion/Config/ModelOptions.cs
@@ -21,6 +21,7 @@ public class ModelOptions : IModelOptions, IOnnxModelSetConfig
         public TokenizerType TokenizerType { get; set; }
         public int SampleSize { get; set; } = 512;
         public float ScaleFactor { get; set; }
+        public ModelType ModelType { get; set; }
         public DiffuserPipelineType PipelineType { get; set; }
         public List<DiffuserType> Diffusers { get; set; } = new List<DiffuserType>();
 
diff --git a/OnnxStack.StableDiffusion/Config/SchedulerOptions.cs b/OnnxStack.StableDiffusion/Config/SchedulerOptions.cs
@@ -87,6 +87,9 @@ public record SchedulerOptions
 
         public int OriginalInferenceSteps { get; set; } = 50;
 
+        public float AestheticScore { get; set; } = 6f;
+        public float AestheticNegativeScore { get; set; } = 2.5f;
+
         public bool IsKarrasScheduler
         {
             get
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/InpaintLegacyDiffuser.cs
@@ -71,6 +71,9 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                 // Get Model metadata
                 var metadata = _onnxModelService.GetModelMetadata(modelOptions, OnnxModelType.Unet);
 
+                // Get Time ids
+                var addTimeIds = GetAddTimeIds(modelOptions, schedulerOptions, performGuidance);
+
                 // Loop though the timesteps
                 var step = 0;
                 foreach (var timestep in timesteps)
@@ -83,7 +86,6 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     var inputLatent = performGuidance ? latents.Repeat(2) : latents;
                     var inputTensor = scheduler.ScaleInput(inputLatent, timestep);
                     var timestepTensor = CreateTimestepTensor(timestep);
-                    var addTimeIds = GetAddTimeIds(schedulerOptions, performGuidance);
 
                     var outputChannels = performGuidance ? 2 : 1;
                     var outputDimension = schedulerOptions.GetScaledDimension(outputChannels);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/StableDiffusionXLDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/StableDiffusionXLDiffuser.cs
@@ -59,6 +59,9 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                 // Get Model metadata
                 var metadata = _onnxModelService.GetModelMetadata(modelOptions, OnnxModelType.Unet);
 
+                // Get Time ids
+                var addTimeIds = GetAddTimeIds(modelOptions, schedulerOptions, performGuidance);
+
                 // Loop though the timesteps
                 var step = 0;
                 foreach (var timestep in timesteps)
@@ -71,7 +74,6 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     var inputLatent = performGuidance ? latents.Repeat(2) : latents;
                     var inputTensor = scheduler.ScaleInput(inputLatent, timestep);
                     var timestepTensor = CreateTimestepTensor(timestep);
-                    var addTimeIds = GetAddTimeIds(schedulerOptions, performGuidance);
 
                     var outputChannels = performGuidance ? 2 : 1;
                     var outputDimension = schedulerOptions.GetScaledDimension(outputChannels);
@@ -113,19 +115,27 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
         /// </summary>
         /// <param name="schedulerOptions">The scheduler options.</param>
         /// <returns></returns>
-        protected DenseTensor<float> GetAddTimeIds(SchedulerOptions schedulerOptions, bool performGuidance)
+        protected DenseTensor<float> GetAddTimeIds(IModelOptions model, SchedulerOptions schedulerOptions, bool performGuidance)
         {
-            var addTimeIds = new float[]
+            float[] result;
+            if (model.ModelType == ModelType.Refiner)
             {
-                schedulerOptions.Height, schedulerOptions.Width, //original_size
-                0, 0, //crops_coords_top_left
-                schedulerOptions.Height, schedulerOptions.Width //negative_target_size
-            };
-            var result = TensorHelper.CreateTensor(addTimeIds, new[] { 1, addTimeIds.Length });
-            if (performGuidance)
-                return result.Repeat(2);
+                //original_size + crops_coords_top_left + aesthetic_score
+                //original_size + crops_coords_top_left + negative_aesthetic_score
+                result = !performGuidance
+                    ? new float[] { schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.AestheticScore }
+                    : new float[] { schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.AestheticNegativeScore, schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.AestheticScore };
+            }
+            else
+            {
+                //original_size + crops_coords_top_left + target_size
+                //original_size + crops_coords_top_left + negative_target_size
+                result = !performGuidance
+                      ? new float[] { schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.Height, schedulerOptions.Width }
+                      : new float[] { schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.Height, schedulerOptions.Width, schedulerOptions.Height, schedulerOptions.Width, 0, 0, schedulerOptions.Height, schedulerOptions.Width };
+            }
 
-            return result;
+            return TensorHelper.CreateTensor(result, new[] { 1, result.Length });
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Enums/ModelType.cs b/OnnxStack.StableDiffusion/Enums/ModelType.cs
@@ -0,0 +1,8 @@
+﻿namespace OnnxStack.StableDiffusion.Enums
+{
+    public enum ModelType
+    {
+        Base = 0,
+        Refiner = 1
+    }
+}
diff --git a/OnnxStack.UI/Models/ModelConfigTemplate.cs b/OnnxStack.UI/Models/ModelConfigTemplate.cs
@@ -23,6 +23,7 @@ public class ModelConfigTemplate
         public List<DiffuserType> Diffusers { get; set; } = new List<DiffuserType>();
         public List<string> ModelFiles { get; set; } = new List<string>();
         public List<string> Images { get; set; } = new List<string>();
+        public ModelType ModelType { get; set; }
     }
 
 }
diff --git a/OnnxStack.UI/Models/ModelSetViewModel.cs b/OnnxStack.UI/Models/ModelSetViewModel.cs
@@ -39,6 +39,7 @@ public class ModelSetViewModel : INotifyPropertyChanged
         private int _dualEmbeddingsLength;
         private TokenizerType _tokenizerType;
         private int _sampleSize;
+        private ModelType _modelType;
 
         public string Name
         {
@@ -215,6 +216,11 @@ public bool HasChanged
             set { _hasChanged = value; NotifyPropertyChanged(); }
         }
 
+        public ModelType ModelType
+        {
+            get { return _modelType; }
+            set { _modelType = value; NotifyPropertyChanged(); }
+        }
 
         public IEnumerable<DiffuserType> GetDiffusers()
         {
diff --git a/OnnxStack.UI/Views/ModelView.xaml.cs b/OnnxStack.UI/Views/ModelView.xaml.cs
@@ -944,6 +944,7 @@ private ModelSetViewModel CreateViewModel(ModelConfigTemplate modelTemplate)
                 ScaleFactor = modelTemplate.ScaleFactor,
                 TokenizerLimit = modelTemplate.TokenizerLimit,
                 SampleSize = modelTemplate.SampleSize,
+                ModelType = modelTemplate.ModelType,
                 PipelineType = modelTemplate.PipelineType,
                 EnableTextToImage = modelTemplate.Diffusers.Contains(DiffuserType.TextToImage),
                 EnableImageToImage = modelTemplate.Diffusers.Contains(DiffuserType.ImageToImage),
@@ -969,6 +970,7 @@ private ModelSetViewModel CreateViewModel(ModelConfigTemplate modelTemplate)
                     ModelFiles = modelTemplate.ModelFiles.ToList(),
                     Repository = modelTemplate.Repository,
                     SampleSize = modelTemplate.SampleSize,
+                    ModelType = modelTemplate.ModelType,
                     Status = ModelTemplateStatus.Installed
                 }
             };
@@ -1001,6 +1003,7 @@ private ModelSetViewModel CreateViewModel(ModelOptions modelOptions)
                 PadTokenId = modelOptions.PadTokenId,
                 ScaleFactor = modelOptions.ScaleFactor,
                 SampleSize = modelOptions.SampleSize,
+                ModelType = modelOptions.ModelType,
                 TokenizerLimit = modelOptions.TokenizerLimit,
                 PipelineType = modelOptions.PipelineType,
                 EnableTextToImage = modelOptions.Diffusers.Contains(DiffuserType.TextToImage),
@@ -1036,6 +1039,7 @@ private ModelSetViewModel CreateViewModel(ModelOptions modelOptions)
                     TokenizerLength = modelOptions.TokenizerLength,
                     Tokenizer2Length = modelOptions.Tokenizer2Length,
                     SampleSize = modelOptions.SampleSize,
+                    ModelType = modelOptions.ModelType,
                     Description = "",
                     Diffusers = modelOptions.Diffusers,
                     ImageIcon = "",
@@ -1074,6 +1078,7 @@ private ModelOptions CreateModelOptions(ModelSetViewModel editModel)
                 PipelineType = editModel.PipelineType,
                 Diffusers = new List<DiffuserType>(editModel.GetDiffusers()),
                 SampleSize = editModel.SampleSize,
+                ModelType = editModel.ModelType,
                 ModelConfigurations = new List<OnnxModelSessionConfig>(editModel.ModelFiles.Select(x => new OnnxModelSessionConfig
                 {
                     Type = x.Type,

Original file line number	Diff line number	Diff line change
`@@ -87,6 +87,9 @@ public record SchedulerOptions`
`87`	`87`
`88`	`88`	`public int OriginalInferenceSteps { get; set; } = 50;`
`89`	`89`
	`90`	`+ public float AestheticScore { get; set; } = 6f;`
	`91`	`+ public float AestheticNegativeScore { get; set; } = 2.5f;`
	`92`	`+`
`90`	`93`	`public bool IsKarrasScheduler`
`91`	`94`	`{`
`92`	`95`	`get`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@ public class ModelConfigTemplate`
`23`	`23`	`public List<DiffuserType> Diffusers { get; set; } = new List<DiffuserType>();`
`24`	`24`	`public List<string> ModelFiles { get; set; } = new List<string>();`
`25`	`25`	`public List<string> Images { get; set; } = new List<string>();`
	`26`	`+ public ModelType ModelType { get; set; }`
`26`	`27`	`}`
`27`	`28`
`28`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ public class ModelSetViewModel : INotifyPropertyChanged`
`39`	`39`	`private int _dualEmbeddingsLength;`
`40`	`40`	`private TokenizerType _tokenizerType;`
`41`	`41`	`private int _sampleSize;`
	`42`	`+ private ModelType _modelType;`
`42`	`43`
`43`	`44`	`public string Name`
`44`	`45`	`{`
`@@ -215,6 +216,11 @@ public bool HasChanged`
`215`	`216`	`set { _hasChanged = value; NotifyPropertyChanged(); }`
`216`	`217`	`}`
`217`	`218`
	`219`	`+ public ModelType ModelType`
	`220`	`+ {`
	`221`	`+ get { return _modelType; }`
	`222`	`+ set { _modelType = value; NotifyPropertyChanged(); }`
	`223`	`+ }`
`218`	`224`
`219`	`225`	`public IEnumerable<DiffuserType> GetDiffusers()`
`220`	`226`	`{`