From ae13c871c89b4c8669539af2d1e4c21c908fe1cb Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Tue, 19 Dec 2023 00:48:08 -0500
Subject: [PATCH] Add model index key length check and skip

---
 CHANGELOG.md                                  |  1 +
 .../Services/ModelIndexService.cs             | 42 ++++++++++---------
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a94c5738..4d546cf4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - Fixed denoise strength in Inference Text to Image
 - Fixed PathTooLongException for IPAdapter folders when using ComfyUI in Symlink mode
 - Fixed configs and symlinks not being cleaned up when switched to the opposite mode
+- Fixed model indexing stopping when encountering paths longer than 1021 bytes in length
 
 ## v2.7.3
 ### Added
diff --git a/StabilityMatrix.Core/Services/ModelIndexService.cs b/StabilityMatrix.Core/Services/ModelIndexService.cs
index dec7eeb9..9733f126 100644
--- a/StabilityMatrix.Core/Services/ModelIndexService.cs
+++ b/StabilityMatrix.Core/Services/ModelIndexService.cs
@@ -1,4 +1,5 @@
 using System.Diagnostics;
+using System.Text;
 using AsyncAwaitBestPractices;
 using Microsoft.Extensions.Logging;
 using StabilityMatrix.Core.Attributes;
@@ -18,8 +19,7 @@ public class ModelIndexService : IModelIndexService
     private readonly ILiteDbContext liteDbContext;
     private readonly ISettingsManager settingsManager;
 
-    public Dictionary<SharedFolderType, List<LocalModelFile>> ModelIndex { get; private set; } =
-        new();
+    public Dictionary<SharedFolderType, List<LocalModelFile>> ModelIndex { get; private set; } = new();
 
     public ModelIndexService(
         ILogger<ModelIndexService> logger,
@@ -48,7 +48,8 @@ public class ModelIndexService : IModelIndexService
     /// <inheritdoc />
     public async Task<IReadOnlyList<LocalModelFile>> GetModelsOfType(SharedFolderType type)
     {
-        return await liteDbContext.LocalModelFiles
+        return await liteDbContext
+            .LocalModelFiles
             .Query()
             .Where(m => m.SharedFolderType == type)
             .ToArrayAsync()
@@ -87,7 +88,8 @@ public class ModelIndexService : IModelIndexService
 
         var newIndex = new Dictionary<SharedFolderType, List<LocalModelFile>>();
         foreach (
-            var file in modelsDir.Info
+            var file in modelsDir
+                .Info
                 .EnumerateFiles("*.*", SearchOption.AllDirectories)
                 .Select(info => new FilePath(info))
         )
@@ -111,16 +113,22 @@ public class ModelIndexService : IModelIndexService
                 continue;
             }
 
-            var localModel = new LocalModelFile
+            // Since RelativePath is the database key, for LiteDB this is limited to 1021 bytes
+            if (Encoding.Unicode.GetByteCount(relativePath) is var byteCount and > 1021)
             {
-                RelativePath = relativePath,
-                SharedFolderType = sharedFolderType,
-            };
+                logger.LogWarning(
+                    "Skipping model {Path} because it's path is too long ({Length} bytes)",
+                    relativePath,
+                    byteCount
+                );
+
+                continue;
+            }
+
+            var localModel = new LocalModelFile { RelativePath = relativePath, SharedFolderType = sharedFolderType, };
 
             // Try to find a connected model info
-            var jsonPath = file.Directory!.JoinFile(
-                new FilePath($"{file.NameWithoutExtension}.cm-info.json")
-            );
+            var jsonPath = file.Directory!.JoinFile(new FilePath($"{file.NameWithoutExtension}.cm-info.json"));
 
             if (jsonPath.Exists)
             {
@@ -132,18 +140,14 @@ public class ModelIndexService : IModelIndexService
             }
 
             // Try to find a preview image
-            var previewImagePath = LocalModelFile.SupportedImageExtensions
-                .Select(
-                    ext => file.Directory!.JoinFile($"{file.NameWithoutExtension}.preview{ext}")
-                )
+            var previewImagePath = LocalModelFile
+                .SupportedImageExtensions
+                .Select(ext => file.Directory!.JoinFile($"{file.NameWithoutExtension}.preview{ext}"))
                 .FirstOrDefault(path => path.Exists);
 
             if (previewImagePath != null)
             {
-                localModel.PreviewImageRelativePath = Path.GetRelativePath(
-                    modelsDir,
-                    previewImagePath
-                );
+                localModel.PreviewImageRelativePath = Path.GetRelativePath(modelsDir, previewImagePath);
             }
 
             // Insert into database