Add custom node model folders + show (#102)

* Revert "Remove custom node search path. (#95)" This reverts commit e09478b. * Add more custom node model overrides.
Comfy-Org · Oct 17, 2024 · 20d79a1 · 20d79a1
1 parent 36ad2f5
commit 20d79a1
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 21 deletions.
diff --git a/src/config/extra_model_config.ts b/src/config/extra_model_config.ts
@@ -30,7 +30,22 @@ const commonPaths = {
   upscale_models: 'models/upscale_models/',
   vae: 'models/vae/',
   vae_approx: 'models/vae_approx/',
-  //custom_nodes: 'custom_nodes/', TODO(robinhuang): https://github.com/Comfy-Org/electron/issues/94
+  // TODO(robinhuang): Remove when we have a better way to specify base model paths.
+  animatediff_models: 'models/animatediff_models/',
+  animatediff_motion_lora: 'models/animatediff_motion_lora/',
+  animatediff_video_formats: 'models/animatediff_video_formats/',
+  ipadater: 'models/ipadater/',
+  liveportrait: 'models/liveportrait/',
+  insightface: 'models/insightface/',
+  layerstyle: 'models/layerstyle/',
+  LLM: 'models/LLM/',
+  Joy_caption: 'models/Joy_caption/',
+  sams: 'models/sams/',
+  blip: 'models/blip/',
+  CogVideo: 'models/CogVideo/',
+  xlabs: 'models/xlabs/',
+  // End custom node model directories.
+  custom_nodes: 'custom_nodes/',
 };
 
 const configTemplates: Record<string, ModelPaths> = {

diff --git a/src/main.ts b/src/main.ts
@@ -173,7 +173,7 @@ if (!gotTheLock) {
       });
       await handleFirstTimeSetup();
       const { appResourcesPath, pythonInstallPath, modelConfigPath, basePath } = await determineResourcesPaths();
-      SetupTray(mainWindow, basePath, () => {
+      SetupTray(mainWindow, basePath, modelConfigPath, () => {
         log.info('Resetting install location');
         fs.rmSync(modelConfigPath);
         restartApp();
@@ -663,9 +663,8 @@ function isComfyUIDirectory(directory: string): boolean {
   return requiredSubdirs.every((subdir) => fs.existsSync(path.join(directory, subdir)));
 }
 
-type DirectoryStructure = (string | [string, string[]])[];
+type DirectoryStructure = (string | DirectoryStructure)[];
 
-// Create directories needed by ComfyUI in the user's data directory.
 function createComfyDirectories(localComfyDirectory: string): void {
   log.info(`Creating ComfyUI directories in ${localComfyDirectory}`);
 
@@ -694,30 +693,54 @@ function createComfyDirectories(localComfyDirectory: string): void {
         'upscale_models',
         'vae',
         'vae_approx',
+
+        // TODO(robinhuang): Remove when we have a better way to specify base model paths.
+        'animatediff_models',
+        'animatediff_motion_lora',
+        'animatediff_video_formats',
+        'liveportrait',
+        ['insightface', ['buffalo_1']],
+        ['blip', ['checkpoints']],
+        'CogVideo',
+        ['xlabs', ['loras', 'controlnets']],
+        'layerstyle',
+        'LLM',
+        'Joy_caption',
       ],
     ],
   ];
-  createDirIfNotExists(localComfyDirectory);
-
-  directories.forEach((dir: string | [string, string[]]) => {
-    if (Array.isArray(dir)) {
-      const [mainDir, subDirs] = dir;
-      const mainDirPath: string = path.join(localComfyDirectory, mainDir);
-      createDirIfNotExists(mainDirPath);
-      subDirs.forEach((subDir: string) => {
-        const subDirPath: string = path.join(mainDirPath, subDir);
-        createDirIfNotExists(subDirPath);
-      });
-    } else {
-      const dirPath: string = path.join(localComfyDirectory, dir);
-      createDirIfNotExists(dirPath);
-    }
-  });
+  try {
+    createNestedDirectories(localComfyDirectory, directories);
+  } catch (error) {
+    log.error(`Failed to create ComfyUI directories: ${error}`);
+  }
 
   const userSettingsPath = path.join(localComfyDirectory, 'user', 'default');
   createComfyConfigFile(userSettingsPath, true);
 }
 
+function createNestedDirectories(basePath: string, structure: DirectoryStructure): void {
+  structure.forEach((item) => {
+    if (typeof item === 'string') {
+      const dirPath = path.join(basePath, item);
+      createDirIfNotExists(dirPath);
+    } else if (Array.isArray(item) && item.length === 2) {
+      const [dirName, subDirs] = item;
+      if (typeof dirName === 'string') {
+        const newBasePath = path.join(basePath, dirName);
+        createDirIfNotExists(newBasePath);
+        if (Array.isArray(subDirs)) {
+          createNestedDirectories(newBasePath, subDirs);
+        }
+      } else {
+        log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`);
+      }
+    } else {
+      log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`);
+    }
+  });
+}
+
 /**
  * Create a directory if not exists
  * @param dirPath

diff --git a/src/tray.ts b/src/tray.ts
@@ -1,7 +1,12 @@
 import { Tray, Menu, BrowserWindow, app, shell } from 'electron';
 import path from 'path';
 
-export function SetupTray(mainView: BrowserWindow, basePath: string, reinstall: () => void): Tray {
+export function SetupTray(
+  mainView: BrowserWindow,
+  basePath: string,
+  modelConfigPath: string,
+  reinstall: () => void
+): Tray {
   // Set icon for the tray
   // I think there is a way to packaged the icon in so you don't need to reference resourcesPath
   const trayImage = path.join(
@@ -68,6 +73,10 @@ export function SetupTray(mainView: BrowserWindow, basePath: string, reinstall:
       label: 'Open Custom Nodes Folder',
       click: () => shell.openPath(path.join(basePath, 'custom_nodes')),
     },
+    {
+      label: 'Open Model Config',
+      click: () => shell.openPath(modelConfigPath),
+    },
     {
       label: 'Open Logs Folder',
       click: () => shell.openPath(app.getPath('logs')),