Read base path from extra_model_config and start comfy with it. (#83)

* Bug fix. * Read base path from existing extra_config_yaml. * Prettier.
Comfy-Org · Oct 15, 2024 · b704d1f · b704d1f
1 parent 4c92a7a
commit b704d1f
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 63 deletions.
diff --git a/package.json b/package.json
@@ -3,7 +3,7 @@
   "productName": "ComfyUI",
   "repository": "github:comfy-org/electron",
   "copyright": "Copyright © 2024 Comfy Org",
-  "version": "0.1.26",
+  "version": "0.1.27",
   "description": "The best modular GUI to run AI diffusion models.",
   "main": ".vite/build/main.js",
   "packageManager": "[email protected]",

diff --git a/src/config/extra_model_config.ts b/src/config/extra_model_config.ts
@@ -1,7 +1,6 @@
 import * as fsPromises from 'node:fs/promises';
-import path from 'path';
 import log from 'electron-log/main';
-import { stringify } from 'yaml';
+import { stringify, parse } from 'yaml';
 
 interface ModelPaths {
   comfyui: {
@@ -85,3 +84,24 @@ export async function createModelConfigFiles(extraModelConfigPath: string, custo
     return false;
   }
 }
+
+export async function readBasePathFromConfig(configPath: string): Promise<string | null> {
+  try {
+    const fileContent = await fsPromises.readFile(configPath, 'utf8');
+    const config = parse(fileContent);
+
+    if (config && config.comfyui && config.comfyui.base_path) {
+      return config.comfyui.base_path;
+    } else {
+      log.warn(`No base_path found in ${configPath}`);
+      return null;
+    }
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
+      log.info(`Config file not found at ${configPath}`);
+    } else {
+      log.error(`Error reading config file ${configPath}:`, error);
+    }
+    return null;
+  }
+}
diff --git a/src/main.ts b/src/main.ts
@@ -11,7 +11,6 @@ import {
   IPCChannel,
   SENTRY_URL_ENDPOINT,
 } from './constants';
-import dotenv from 'dotenv';
 import { app, BrowserWindow, dialog, screen, ipcMain, Menu, MenuItem } from 'electron';
 import tar from 'tar';
 import log from 'electron-log/main';
@@ -20,27 +19,31 @@ import Store from 'electron-store';
 import { updateElectronApp, UpdateSourceType } from 'update-electron-app';
 import * as net from 'net';
 import { graphics } from 'systeminformation';
-import { createModelConfigFiles } from './config/extra_model_config';
+import { createModelConfigFiles, readBasePathFromConfig } from './config/extra_model_config';
 
 let comfyServerProcess: ChildProcess | null = null;
 const host = '127.0.0.1';
 let port = 8188;
 let mainWindow: BrowserWindow | null;
+let store: Store<StoreType> | null;
 const messageQueue: Array<any> = []; // Stores mesaages before renderer is ready.
 
 import { StoreType } from './store';
 log.initialize();
 // Handle creating/removing shortcuts on Windows when installing/uninstalling.
 // Run this as early in the main process as possible.
-if (require('electron-squirrel-startup')) app.quit();
-
-const store = new Store<StoreType>();
+if (require('electron-squirrel-startup')) {
+  log.info('App already being set up by squirrel. Exiting...');
+  app.quit();
+}
 
 const gotTheLock = app.requestSingleInstanceLock();
 
 if (!gotTheLock) {
+  log.info('App already running. Exiting...');
   app.quit();
 } else {
+  store = new Store<StoreType>();
   app.on('second-instance', (event, commandLine, workingDirectory, additionalData) => {
     log.info('Received second instance message!');
     log.info(additionalData);
@@ -137,7 +140,7 @@ if (!gotTheLock) {
         });
       });
       await handleFirstTimeSetup();
-      const { userResourcesPath, appResourcesPath, pythonInstallPath, modelConfigPath } =
+      const { userResourcesPath, appResourcesPath, pythonInstallPath, modelConfigPath, basePath } =
         await determineResourcesPaths();
       SetupTray(mainWindow, userResourcesPath);
       port = await findAvailablePort(8000, 9999).catch((err) => {
@@ -148,7 +151,7 @@ if (!gotTheLock) {
       sendProgressUpdate('Setting up Python Environment...');
       const pythonInterpreterPath = await setupPythonEnvironment(appResourcesPath, pythonInstallPath);
       sendProgressUpdate('Starting Comfy Server...');
-      await launchPythonServer(pythonInterpreterPath, appResourcesPath, userResourcesPath, modelConfigPath);
+      await launchPythonServer(pythonInterpreterPath, appResourcesPath, modelConfigPath, basePath);
       updateElectronApp({
         updateSource: {
           type: UpdateSourceType.StaticStorage,
@@ -202,6 +205,7 @@ if (!gotTheLock) {
   app.on('window-all-closed', () => {
     log.info('Window all closed');
     if (process.platform !== 'darwin') {
+      log.info('Quitting ComfyUI because window all closed');
       app.quit();
     }
   });
@@ -374,8 +378,8 @@ let spawnServerTimeout: NodeJS.Timeout = null;
 const launchPythonServer = async (
   pythonInterpreterPath: string,
   appResourcesPath: string,
-  userResourcesPath: string,
-  modelConfigPath: string
+  modelConfigPath: string,
+  basePath: string
 ) => {
   const isServerRunning = await isComfyServerReady(host, port);
   if (isServerRunning) {
@@ -388,9 +392,9 @@ const launchPythonServer = async (
 
   return new Promise<void>(async (resolve, reject) => {
     const scriptPath = path.join(appResourcesPath, 'ComfyUI', 'main.py');
-    const userDirectoryPath = path.join(userResourcesPath, 'user');
-    const inputDirectoryPath = path.join(userResourcesPath, 'input');
-    const outputDirectoryPath = path.join(userResourcesPath, 'output');
+    const userDirectoryPath = path.join(basePath, 'user');
+    const inputDirectoryPath = path.join(basePath, 'input');
+    const outputDirectoryPath = path.join(basePath, 'output');
     const comfyMainCmd = [
       scriptPath,
       '--user-directory',
@@ -544,16 +548,6 @@ const spawnPython = (
         mainWindow.webContents.send(IPC_CHANNELS.LOG_MESSAGE, message);
       }
     });
-
-    const signalHandler = (signal: NodeJS.Signals) => {
-      log.warn(`Received ${signal}, terminating Python process`);
-      if (!pythonProcess.killed) {
-        pythonProcess.kill(signal); // Send the signal to the Python process
-      }
-    };
-
-    process.on('SIGINT', signalHandler);
-    process.on('SIGTERM', signalHandler);
   }
 
   return pythonProcess;
@@ -569,15 +563,8 @@ const spawnPythonAsync = (
     log.info(`Spawning python process with command: ${cmd.join(' ')} in directory: ${cwd}`);
     const pythonProcess: ChildProcess = spawn(pythonInterpreterPath, cmd, { cwd });
 
-    let timeoutId: NodeJS.Timeout | null = null;
-
     const cleanup = () => {
-      if (timeoutId) clearTimeout(timeoutId);
       pythonProcess.removeAllListeners();
-      if (!pythonProcess.killed) {
-        pythonProcess.kill();
-      }
-      process.exit();
     };
 
     if (options.stdx) {
@@ -609,18 +596,6 @@ const spawnPythonAsync = (
       log.error(`Failed to start Python process: ${err}`);
       reject(err);
     });
-
-    const signalHandler = (signal: NodeJS.Signals) => {
-      log.warn(`Received ${signal}, terminating Python process`);
-      cleanup();
-      if (!pythonProcess.killed) {
-        pythonProcess.kill(signal);
-      }
-      process.exit();
-    };
-
-    process.on('SIGINT', signalHandler);
-    process.on('SIGTERM', signalHandler);
   });
 };
 
@@ -865,15 +840,18 @@ async function determineResourcesPaths(): Promise<{
   pythonInstallPath: string;
   appResourcesPath: string;
   modelConfigPath: string;
+  basePath: string | null;
 }> {
   const modelConfigPath = path.join(app.getPath('userData'), 'extra_models_config.yaml');
+  const basePath = await readBasePathFromConfig(modelConfigPath);
   if (!app.isPackaged) {
     return {
       // development: install python to in-tree assets dir
       userResourcesPath: path.join(app.getAppPath(), 'assets'),
       pythonInstallPath: path.join(app.getAppPath(), 'assets'),
       appResourcesPath: path.join(app.getAppPath(), 'assets'),
       modelConfigPath: modelConfigPath,
+      basePath: basePath,
     };
   }
 
@@ -891,6 +869,7 @@ async function determineResourcesPaths(): Promise<{
     pythonInstallPath: defaultPythonInstallPath,
     appResourcesPath: appResourcePath,
     modelConfigPath: modelConfigPath,
+    basePath: basePath,
   };
 }
 

diff --git a/src/renderer/index.tsx b/src/renderer/index.tsx
@@ -49,7 +49,7 @@ const Home: React.FC = () => {
       <div style={bodyStyle}>
         <FirstTimeSetup onComplete={() => setShowSetup(false)} />
       </div>
-    )
+    );
   }
 
   return (

diff --git a/src/renderer/screens/FirstTimeSetup.tsx b/src/renderer/screens/FirstTimeSetup.tsx
@@ -38,34 +38,44 @@ const FirstTimeSetup: React.FC<FirstTimeSetupProps> = ({ onComplete }) => {
   return (
     <div style={styles.container}>
       <h1 style={styles.title}>Install ComfyUI</h1>
-      {!selectedPath && <p style={styles.description}>
-        Please select a directory for where ComfyUI will store models, outputs, etc. If you already have a ComfyUI
-        setup, you can select that to reuse the model files.
-      </p>}
-      {!selectedPath && <button onClick={handleDirectorySelect} style={styles.selectButton}>
-        Select Directory
-      </button>}
+      {!selectedPath && (
+        <p style={styles.description}>
+          Please select a directory for where ComfyUI will store models, outputs, etc. If you already have a ComfyUI
+          setup, you can select that to reuse the model files.
+        </p>
+      )}
+      {!selectedPath && (
+        <button onClick={handleDirectorySelect} style={styles.selectButton}>
+          Select Directory
+        </button>
+      )}
       {selectedPath && (
         <div style={styles.pathDisplay}>
-          <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" strokeWidth={1.5} stroke="currentColor" className="size-6" style={styles.folderIcon}>
-            <path strokeLinecap="round" strokeLinejoin="round" d="M2.25 12.75V12A2.25 2.25 0 0 1 4.5 9.75h15A2.25 2.25 0 0 1 21.75 12v.75m-8.69-6.44-2.12-2.12a1.5 1.5 0 0 0-1.061-.44H4.5A2.25 2.25 0 0 0 2.25 6v12a2.25 2.25 0 0 0 2.25 2.25h15A2.25 2.25 0 0 0 21.75 18V9a2.25 2.25 0 0 0-2.25-2.25h-5.379a1.5 1.5 0 0 1-1.06-.44Z" />
+          <svg
+            xmlns="http://www.w3.org/2000/svg"
+            fill="none"
+            viewBox="0 0 24 24"
+            strokeWidth={1.5}
+            stroke="currentColor"
+            className="size-6"
+            style={styles.folderIcon}
+          >
+            <path
+              strokeLinecap="round"
+              strokeLinejoin="round"
+              d="M2.25 12.75V12A2.25 2.25 0 0 1 4.5 9.75h15A2.25 2.25 0 0 1 21.75 12v.75m-8.69-6.44-2.12-2.12a1.5 1.5 0 0 0-1.061-.44H4.5A2.25 2.25 0 0 0 2.25 6v12a2.25 2.25 0 0 0 2.25 2.25h15A2.25 2.25 0 0 0 21.75 18V9a2.25 2.25 0 0 0-2.25-2.25h-5.379a1.5 1.5 0 0 1-1.06-.44Z"
+            />
           </svg>
 
           <p style={styles.pathText}>{selectedPath}</p>
         </div>
       )}
       {selectedPath && (
         <div style={styles.buttonContainer}>
-          <button
-            onClick={() => setSelectedPath('')}
-            style={styles.changePathButton}
-          >
+          <button onClick={() => setSelectedPath('')} style={styles.changePathButton}>
             Back
           </button>
-          <button
-            onClick={handleInstall}
-            style={styles.installButton}
-          >
+          <button onClick={handleInstall} style={styles.installButton}>
             Install
           </button>
         </div>

diff --git a/src/renderer/screens/ProgressOverlay.tsx b/src/renderer/screens/ProgressOverlay.tsx
@@ -93,7 +93,9 @@ function ProgressOverlay(): React.ReactElement {
         <div style={loadingTextStyle} id="loading-text">
           {status}
         </div>
-        <div style={logContainerStyle}>{status !== COMFY_FINISHING_MESSAGE && status !== COMFY_ERROR_MESSAGE && <AnimatedLogDisplay logs={logs} />}</div>
+        <div style={logContainerStyle}>
+          {status !== COMFY_FINISHING_MESSAGE && status !== COMFY_ERROR_MESSAGE && <AnimatedLogDisplay logs={logs} />}
+        </div>
       </div>
     </div>
   );
-Original file line number
+Diff line change
@@ Expand Up / @@ -49,7 +49,7 @@ const Home: React.FC = () => { @@
           <div style={bodyStyle}>
             <FirstTimeSetup onComplete={() => setShowSetup(false)} />
           </div>
-        )
+        );
       }
       return (
@@ Expand Down @@