Merge pull request LykosAI#749 from ionite34/add-flux-types

Added support for downloading/browsing Flux models in the model browser...

Author: mohnjiles

CHANGELOG.md

@@ -8,12 +8,15 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 ## v2.12.0-dev.3
 ### Added
 - Added Settings option "Console: History Size" to adjust the number of lines stored in the console history when running packages. Defaults to 9001 lines.
+#### Model Browser
+- Added AuraFlow & Flux base model types to the CivitAI model browser
 #### Checkpoint Manager
 - Added "New Directory" and "Delete" options to the context menu of the tree view.
 - Added new toggle for drag & drop - when enabled, all selected models will now move together with the dragged model
 - Added "File Size" sorting option
 - Added "Hide Empty Categories" toggle
 - Added "Select All" button to the InfoBar (shown when at least one model is selected)
+- Added "unet" shared model folder for ComfyUI
 ### Changed
 - The "Download Failed" message for model downloads is now persistent until dismissed
 ### Fixed

StabilityMatrix.Avalonia/Assets/hf-packages.json

@@ -922,5 +922,32 @@
     ],
     "Subfolder": "segm",
     "LicenseType": "Apache 2.0"
+  },
+  {
+    "ModelCategory": "Clip",
+    "ModelName": "CLIP L",
+    "RepositoryPath": "comfyanonymous/flux_text_encoders",
+    "Files": [
+      "clip_l.safetensors"
+    ],
+    "LicenseType": "MIT"
+  },
+  {
+    "ModelCategory": "Clip",
+    "ModelName": "T5XXL FP16",
+    "RepositoryPath": "comfyanonymous/flux_text_encoders",
+    "Files": [
+      "t5xxl_fp16.safetensors"
+    ],
+    "LicenseType": "Apache 2.0"
+  },
+  {
+    "ModelCategory": "Clip",
+    "ModelName": "T5XXL FP8 E4M3FN",
+    "RepositoryPath": "comfyanonymous/flux_text_encoders",
+    "Files": [
+      "t5xxl_fp8_e4m3fn.safetensors"
+    ],
+    "LicenseType": "Apache 2.0"
   }
 ]

StabilityMatrix.Avalonia/Models/HuggingFace/HuggingFaceModelType.cs

@@ -12,6 +12,10 @@ public enum HuggingFaceModelType
     [ConvertTo<SharedFolderType>(SharedFolderType.StableDiffusion)]
     BaseModel,
 
+    [Description("CLIP / Text Encoders")]
+    [ConvertTo<SharedFolderType>(SharedFolderType.CLIP)]
+    Clip,
+
     [Description("ControlNets (SD1.5)")]
     [ConvertTo<SharedFolderType>(SharedFolderType.ControlNet)]
     ControlNet,

StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CheckpointBrowserCardViewModel.cs

@@ -253,9 +253,19 @@ private async Task ShowVersionDialog(CivitModel model)
         }
         else
         {
-            var subFolder =
-                viewModel?.SelectedInstallLocation
-                ?? Path.Combine(@"Models", model.Type.ConvertTo<SharedFolderType>().GetStringValue());
+            var sharedFolder = model.Type.ConvertTo<SharedFolderType>().GetStringValue();
+
+            if (
+                model.BaseModelType == CivitBaseModelType.Flux1D.GetStringValue()
+                || model.BaseModelType == CivitBaseModelType.Flux1S.GetStringValue()
+            )
+            {
+                sharedFolder = SharedFolderType.Unet.GetStringValue();
+            }
+
+            var defaultPath = Path.Combine(@"Models", sharedFolder);
+
+            var subFolder = viewModel?.SelectedInstallLocation ?? defaultPath;
             downloadPath = Path.Combine(settingsManager.LibraryDir, subFolder);
         }
 

StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs

@@ -380,6 +380,8 @@ or nameof(SortConnectedModelsFirst)
 
         // make sure a sort happens
         OnPropertyChanged(nameof(SortConnectedModelsFirst));
+        // gotta do this one too i guess
+        OnPropertyChanged(nameof(HideEmptyRootCategories));
     }
 
     public void ClearSearchQuery()

StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs

@@ -306,10 +306,12 @@ private void LoadInstallLocations()
         var installLocations = new ObservableCollection<string>();
 
         var rootModelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory);
-        var downloadDirectory = rootModelsDirectory.JoinDir(
-            SelectedFile?.CivitFile.Type == CivitFileType.VAE
-                ? SharedFolderType.VAE.GetStringValue()
-                : CivitModel.Type.ConvertTo<SharedFolderType>().GetStringValue()
+
+        var downloadDirectory = GetSharedFolderPath(
+            rootModelsDirectory,
+            SelectedFile?.CivitFile.Type,
+            CivitModel.Type,
+            CivitModel.BaseModelType
         );
 
         if (!downloadDirectory.ToString().EndsWith("Unknown"))
@@ -330,4 +332,27 @@ var directory in downloadDirectory.EnumerateDirectories(
         AvailableInstallLocations = installLocations;
         SelectedInstallLocation = installLocations.FirstOrDefault();
     }
+
+    private static DirectoryPath GetSharedFolderPath(
+        DirectoryPath rootModelsDirectory,
+        CivitFileType? fileType,
+        CivitModelType modelType,
+        string? baseModelType
+    )
+    {
+        if (fileType is CivitFileType.VAE)
+        {
+            return rootModelsDirectory.JoinDir(SharedFolderType.VAE.GetStringValue());
+        }
+
+        if (
+            baseModelType == CivitBaseModelType.Flux1D.GetStringValue()
+            || baseModelType == CivitBaseModelType.Flux1S.GetStringValue()
+        )
+        {
+            return rootModelsDirectory.JoinDir(SharedFolderType.Unet.GetStringValue());
+        }
+
+        return rootModelsDirectory.JoinDir(modelType.ConvertTo<SharedFolderType>().GetStringValue());
+    }
 }

StabilityMatrix.Core/Models/Api/CivitBaseModelType.cs

@@ -8,6 +8,15 @@ public enum CivitBaseModelType
 {
     All,
 
+    [StringValue("AuraFlow")]
+    AuraFlow,
+
+    [StringValue("Flux.1 S")]
+    Flux1S,
+
+    [StringValue("Flux.1 D")]
+    Flux1D,
+
     [StringValue("PixArt a")]
     PixArtA,
 

StabilityMatrix.Core/Models/Database/LocalModelFile.cs

@@ -201,7 +201,8 @@ public IEnumerable<string> GetDeleteFullPaths(string rootModelDirectory)
         ".pt",
         ".ckpt",
         ".pth",
-        ".bin"
+        ".bin",
+        ".sft"
     ];
     public static readonly HashSet<string> SupportedImageExtensions = [".png", ".jpg", ".jpeg", ".webp"];
     public static readonly HashSet<string> SupportedMetadataExtensions = [".json"];

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -69,7 +69,8 @@ IPrerequisiteHelper prerequisiteHelper
             [SharedFolderType.T2IAdapter] = ["models/controlnet/T2IAdapter"],
             [SharedFolderType.PromptExpansion] = ["models/prompt_expansion"],
             [SharedFolderType.Ultralytics] = ["models/ultralytics"],
-            [SharedFolderType.Sams] = ["models/sams"]
+            [SharedFolderType.Sams] = ["models/sams"],
+            [SharedFolderType.Unet] = ["models/unet"]
         };
 
     public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders =>
@@ -372,6 +373,7 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
             nodeValue.Children["ultralytics_bbox"] = Path.Combine(modelsDir, "Ultralytics", "bbox");
             nodeValue.Children["ultralytics_segm"] = Path.Combine(modelsDir, "Ultralytics", "segm");
             nodeValue.Children["sams"] = Path.Combine(modelsDir, "Sams");
+            nodeValue.Children["unet"] = Path.Combine(modelsDir, "unet");
         }
         else
         {
@@ -414,7 +416,8 @@ private async Task SetupModelFoldersConfig(DirectoryPath installDirectory)
                     { "ultralytics", Path.Combine(modelsDir, "Ultralytics") },
                     { "ultralytics_bbox", Path.Combine(modelsDir, "Ultralytics", "bbox") },
                     { "ultralytics_segm", Path.Combine(modelsDir, "Ultralytics", "segm") },
-                    { "sams", Path.Combine(modelsDir, "Sams") }
+                    { "sams", Path.Combine(modelsDir, "Sams") },
+                    { "unet", Path.Combine(modelsDir, "unet") }
                 }
             );
         }

StabilityMatrix.Core/Models/SharedFolderType.cs

@@ -43,6 +43,6 @@ public enum SharedFolderType : ulong
     SVD = 1 << 27,
     Ultralytics = 1 << 28,
     Sams = 1 << 29,
-
-    PromptExpansion = 1 << 30
+    PromptExpansion = 1 << 30,
+    Unet = 1ul << 31,
 }