Merge pull request #833 from jeremydk/comfyui_unet

Quickly adding unet path for ComfyUI
LykosAI · Aug 20, 2024 · dcf8a40 · dcf8a40
2 parents baf7331 + 7d7e7d1
commit dcf8a40
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/StabilityMatrix.Core/Models/Packages/ComfyUI.cs b/StabilityMatrix.Core/Models/Packages/ComfyUI.cs
@@ -67,7 +67,8 @@ IPrerequisiteHelper prerequisiteHelper
             [SharedFolderType.InvokeIpAdapters15] = new[] { "models/ipadapter/sd15" },
             [SharedFolderType.InvokeIpAdaptersXl] = new[] { "models/ipadapter/sdxl" },
             [SharedFolderType.T2IAdapter] = new[] { "models/controlnet/T2IAdapter" },
-            [SharedFolderType.PromptExpansion] = new[] { "models/prompt_expansion" }
+            [SharedFolderType.PromptExpansion] = new[] { "models/prompt_expansion" },
+            [SharedFolderType.Unet] = new[] { "models/unet" }
         };
 
     public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders =>
@@ -373,6 +374,7 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
                 Path.Combine(modelsDir, "InvokeIpAdaptersXl")
             );
             nodeValue.Children["prompt_expansion"] = Path.Combine(modelsDir, "PromptExpansion");
+            nodeValue.Children["unet"] = Path.Combine(modelsDir, "unet");
         }
         else
         {
@@ -411,7 +413,8 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
                             Path.Combine(modelsDir, "InvokeIpAdaptersXl")
                         )
                     },
-                    { "prompt_expansion", Path.Combine(modelsDir, "PromptExpansion") }
+                    { "prompt_expansion", Path.Combine(modelsDir, "PromptExpansion") },
+                    { "unet", Path.Combine(modelsDir, "unet") },
                 }
             );
         }

diff --git a/StabilityMatrix.Core/Models/SharedFolderType.cs b/StabilityMatrix.Core/Models/SharedFolderType.cs
@@ -43,5 +43,6 @@ public enum SharedFolderType
     InvokeClipVision = 1 << 26,
     SVD = 1 << 27,
 
-    PromptExpansion = 1 << 30
+    PromptExpansion = 1 << 30,
+    Unet = 1 << 31
 }