From 660dad42454c54016fa41fc961068471567aca7b Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Thu, 4 Apr 2024 23:22:39 -0400
Subject: [PATCH] Add layer diffuse selection for SD1.5 / SDXL

---
 .../Inference/LayerDiffuseCardViewModel.cs    | 22 +++++++++++++++++--
 .../Api/Comfy/Nodes/ComfyNodeBuilder.cs       |  6 +++++
 .../Models/Inference/LayerDiffuseMode.cs      | 12 +++++++---
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/LayerDiffuseCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/LayerDiffuseCardViewModel.cs
index 04cb8b80..865a4c74 100644
--- a/StabilityMatrix.Avalonia/ViewModels/Inference/LayerDiffuseCardViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Inference/LayerDiffuseCardViewModel.cs
@@ -37,6 +37,24 @@ public partial class LayerDiffuseCardViewModel : LoadableViewModelBase, IComfySt
         if (SelectedMode == LayerDiffuseMode.None)
             return;
 
+        var sdType = SelectedMode switch
+        {
+            LayerDiffuseMode.GenerateForegroundWithTransparencySD15 => "SD15",
+            LayerDiffuseMode.GenerateForegroundWithTransparencySDXL => "SDXL",
+            LayerDiffuseMode.None => throw new ArgumentOutOfRangeException(),
+            _ => throw new ArgumentOutOfRangeException()
+        };
+
+        // Choose config based on mode
+        var config = SelectedMode switch
+        {
+            LayerDiffuseMode.GenerateForegroundWithTransparencySD15
+                => "SD15, Attention Injection, attn_sharing",
+            LayerDiffuseMode.GenerateForegroundWithTransparencySDXL => "SDXL, Conv Injection",
+            LayerDiffuseMode.None => throw new ArgumentOutOfRangeException(),
+            _ => throw new ArgumentOutOfRangeException()
+        };
+
         foreach (var modelConnections in e.Temp.Models.Values)
         {
             var layerDiffuseApply = e.Nodes.AddTypedNode(
@@ -44,7 +62,7 @@ public partial class LayerDiffuseCardViewModel : LoadableViewModelBase, IComfySt
                 {
                     Name = e.Nodes.GetUniqueName($"LayerDiffuseApply_{modelConnections.Name}"),
                     Model = modelConnections.Model,
-                    Config = "SD15, Attention Injection, attn_sharing",
+                    Config = config,
                     Weight = Weight,
                 }
             );
@@ -71,7 +89,7 @@ public partial class LayerDiffuseCardViewModel : LoadableViewModelBase, IComfySt
                     Name = applyArgs.Nodes.GetUniqueName("LayerDiffuseDecode"),
                     Samples = latent,
                     Images = primaryImage,
-                    SdVersion = "SD15",
+                    SdVersion = sdType
                 }
             );
 
diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
index a1b79df5..df14284a 100644
--- a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
+++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs
@@ -391,6 +391,12 @@ public class ComfyNodeBuilder
     {
         public required ModelNodeConnection Model { get; init; }
 
+        /// <summary>
+        /// Available configs:
+        /// <para>SD15, Attention Injection, attn_sharing</para>
+        /// <para>SDXL, Conv Injection</para>
+        /// <para>SDXL, Attention Injection</para>
+        /// </summary>
         public required string Config { get; init; }
 
         [Range(-1d, 3d)]
diff --git a/StabilityMatrix.Core/Models/Inference/LayerDiffuseMode.cs b/StabilityMatrix.Core/Models/Inference/LayerDiffuseMode.cs
index af243113..35621f60 100644
--- a/StabilityMatrix.Core/Models/Inference/LayerDiffuseMode.cs
+++ b/StabilityMatrix.Core/Models/Inference/LayerDiffuseMode.cs
@@ -11,8 +11,14 @@ public enum LayerDiffuseMode
     None,
 
     /// <summary>
-    /// Generate foreground only with transparency.
+    /// Generate foreground only with transparency. SD1.5
     /// </summary>
-    [Display(Name = "Generate Foreground with Transparency")]
-    GenerateForegroundWithTransparency,
+    [Display(Name = "(SD 1.5) Generate Foreground with Transparency")]
+    GenerateForegroundWithTransparencySD15,
+
+    /// <summary>
+    /// Generate foreground only with transparency. SDXL
+    /// </summary>
+    [Display(Name = "(SDXL) Generate Foreground with Transparency")]
+    GenerateForegroundWithTransparencySDXL,
 }