From ae70f7e5f01547504544f9ab8fb172a6aa2783f4 Mon Sep 17 00:00:00 2001
From: Darth Affe <darthaffe@wyrez.org>
Date: Sun, 25 Aug 2024 16:41:19 +0200
Subject: [PATCH] Changed vae to be required for flux models and allow all non
 null inputs

---
 .../Extensions/DiffusionModelBuilderExtension.cs       | 10 +++++-----
 StableDiffusion.NET/Models/Builder/FluxModelBuilder.cs |  4 ++--
 StableDiffusion.NET/Models/Builder/ModelBuilder.cs     |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/StableDiffusion.NET/Models/Builder/Extensions/DiffusionModelBuilderExtension.cs b/StableDiffusion.NET/Models/Builder/Extensions/DiffusionModelBuilderExtension.cs
index 9f08791..e156e30 100644
--- a/StableDiffusion.NET/Models/Builder/Extensions/DiffusionModelBuilderExtension.cs
+++ b/StableDiffusion.NET/Models/Builder/Extensions/DiffusionModelBuilderExtension.cs
@@ -9,7 +9,7 @@ public static class DiffusionModelBuilderExtension
     public static T WithVae<T>(this T builder, string vaePath)
         where T : IDiffusionModelBuilder
     {
-        ArgumentException.ThrowIfNullOrWhiteSpace(vaePath);
+        ArgumentNullException.ThrowIfNull(vaePath);
 
         if (!string.IsNullOrWhiteSpace(builder.Parameter.TaesdPath)) throw new ArgumentException("TAESD is already enabled. VAE and TAESD are mutually exclusive.", nameof(vaePath));
 
@@ -21,7 +21,7 @@ public static class DiffusionModelBuilderExtension
     public static T WithTaesd<T>(this T builder, string taesdPath)
         where T : IDiffusionModelBuilder
     {
-        ArgumentException.ThrowIfNullOrWhiteSpace(taesdPath);
+        ArgumentNullException.ThrowIfNull(taesdPath);
 
         if (!string.IsNullOrWhiteSpace(builder.Parameter.VaePath)) throw new ArgumentException("VAE is already enabled. TAESD and VAE are mutually exclusive.", nameof(taesdPath));
 
@@ -33,7 +33,7 @@ public static class DiffusionModelBuilderExtension
     public static T WithLoraSupport<T>(this T builder, string loraModelDirectory)
         where T : IDiffusionModelBuilder
     {
-        ArgumentException.ThrowIfNullOrWhiteSpace(loraModelDirectory);
+        ArgumentNullException.ThrowIfNull(loraModelDirectory);
 
         builder.Parameter.LoraModelDirectory = loraModelDirectory;
 
@@ -43,7 +43,7 @@ public static class DiffusionModelBuilderExtension
     public static T WithEmbeddingSupport<T>(this T builder, string embeddingsDirectory)
         where T : IDiffusionModelBuilder
     {
-        ArgumentException.ThrowIfNullOrWhiteSpace(embeddingsDirectory);
+        ArgumentNullException.ThrowIfNull(embeddingsDirectory);
 
         builder.Parameter.EmbeddingsDirectory = embeddingsDirectory;
 
@@ -53,7 +53,7 @@ public static class DiffusionModelBuilderExtension
     public static T WithControlNet<T>(this T builder, string controlNetPath)
         where T : IDiffusionModelBuilder
     {
-        ArgumentException.ThrowIfNullOrWhiteSpace(controlNetPath);
+        ArgumentNullException.ThrowIfNull(controlNetPath);
 
         builder.Parameter.ControlNetPath = controlNetPath;
 
diff --git a/StableDiffusion.NET/Models/Builder/FluxModelBuilder.cs b/StableDiffusion.NET/Models/Builder/FluxModelBuilder.cs
index 415774f..367e23b 100644
--- a/StableDiffusion.NET/Models/Builder/FluxModelBuilder.cs
+++ b/StableDiffusion.NET/Models/Builder/FluxModelBuilder.cs
@@ -15,9 +15,9 @@ public sealed class FluxModelBuilder : IDiffusionModelBuilder, IQuantizedModelBu
 
     #region Constructors
 
-    public FluxModelBuilder(string diffusionModelPath, string clipLPath, string t5xxlPath)
+    public FluxModelBuilder(string diffusionModelPath, string clipLPath, string t5xxlPath, string vaePath)
     {
-        Parameter = new DiffusionModelParameter { DiffusionModelType = DiffusionModelType.Flux, DiffusionModelPath = diffusionModelPath, ClipLPath = clipLPath, T5xxlPath = t5xxlPath };
+        Parameter = new DiffusionModelParameter { DiffusionModelType = DiffusionModelType.Flux, DiffusionModelPath = diffusionModelPath, ClipLPath = clipLPath, T5xxlPath = t5xxlPath, VaePath = vaePath };
     }
 
     #endregion
diff --git a/StableDiffusion.NET/Models/Builder/ModelBuilder.cs b/StableDiffusion.NET/Models/Builder/ModelBuilder.cs
index d56dfc7..d7189bf 100644
--- a/StableDiffusion.NET/Models/Builder/ModelBuilder.cs
+++ b/StableDiffusion.NET/Models/Builder/ModelBuilder.cs
@@ -6,6 +6,6 @@ namespace StableDiffusion.NET;
 public static class ModelBuilder
 {
     public static StableDiffusionModelBuilder StableDiffusion(string modelPath) => new(modelPath);
-    public static FluxModelBuilder Flux(string diffusionModelPath, string clipLPath, string t5xxlPath) => new(diffusionModelPath, clipLPath, t5xxlPath);
+    public static FluxModelBuilder Flux(string diffusionModelPath, string clipLPath, string t5xxlPath, string vaePath) => new(diffusionModelPath, clipLPath, t5xxlPath, vaePath);
     public static ESRGANModelBuilder ESRGAN(string modelPath) => new(modelPath);
 }
\ No newline at end of file