From bbeb6913ff630d2903e83fa52d33580e68e0667c Mon Sep 17 00:00:00 2001
From: N00MKRAD <login@nmkd.de>
Date: Sun, 20 Jun 2021 19:33:23 +0200
Subject: [PATCH] Added support for custom models

Dir name must end with "_custom", must be only A-Z, 0-9 and underscores
---
 Code/Main/AiModels.cs | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/Code/Main/AiModels.cs b/Code/Main/AiModels.cs
index 5d8ff36..53c268c 100644
--- a/Code/Main/AiModels.cs
+++ b/Code/Main/AiModels.cs
@@ -5,6 +5,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Text;
+using System.Text.RegularExpressions;
 using System.Threading.Tasks;
 
 namespace Flowframes.Main
@@ -16,9 +17,31 @@ namespace Flowframes.Main
             string pkgPath = Path.Combine(Paths.GetPkgPath(), ai.pkgDir);
             string modelsFile = Path.Combine(pkgPath, "models.json");
             ModelCollection modelCollection = new ModelCollection(ai, modelsFile);
+
+            foreach (string customModel in GetCustomModels(ai))
+            {
+                string name = customModel.Remove("_alpha").Remove("_custom");
+                bool alpha = customModel.Contains("_alpha");
+                modelCollection.models.Add(new ModelCollection.ModelInfo(ai, name, "Custom Model", customModel, alpha, false));
+            }
+
             return modelCollection;
         }
 
+        public static List<string> GetCustomModels(AI ai)
+        {
+            string pkgPath = Path.Combine(Paths.GetPkgPath(), ai.pkgDir);
+            List<string> custModels = new List<string>();
+
+            foreach (DirectoryInfo dir in new DirectoryInfo(pkgPath).GetDirectories())
+            {
+                if (dir.Name.EndsWith("_custom") && Regex.IsMatch(dir.Name, @"^[a-zA-Z0-9_]+$"))
+                    custModels.Add(dir.Name);
+            }
+
+            return custModels;
+        }
+
         public static ModelCollection.ModelInfo GetModelByName(AI ai, string modelName)
         {
             ModelCollection modelCollection = GetModels(ai);