From 75a230adc7ffe2a3bea221a8a9ddd0fe78b69c66 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <hcl@comfy.org>
Date: Tue, 5 Nov 2024 14:43:19 -0500
Subject: [PATCH] Make port/host configurable (#178)

---
 .env_example      |  8 ++++++
 .env_test_example | 10 +++++++
 src/main.ts       | 69 +++++++++++++++++++++++++++++++----------------
 3 files changed, 64 insertions(+), 23 deletions(-)
 create mode 100644 .env_example
 create mode 100644 .env_test_example

diff --git a/.env_example b/.env_example
new file mode 100644
index 00000000..907fead8
--- /dev/null
+++ b/.env_example
@@ -0,0 +1,8 @@
+# The host to use for the ComfyUI server.
+COMFY_HOST=127.0.0.1
+
+# The port to use for the ComfyUI server.
+COMFY_PORT=8188
+
+# Whether to use an external server instead of starting one locally.
+USE_EXTERNAL_SERVER=false
diff --git a/.env_test_example b/.env_test_example
new file mode 100644
index 00000000..94f2d29b
--- /dev/null
+++ b/.env_test_example
@@ -0,0 +1,10 @@
+# Use this file to test with local ComfyUI_frontend dev server.
+
+# The host to use for the ComfyUI server.
+COMFY_HOST=localhost
+
+# The port to use for the ComfyUI server.
+COMFY_PORT=5173
+
+# Whether to use an external server instead of starting one locally.
+USE_EXTERNAL_SERVER=true
\ No newline at end of file
diff --git a/src/main.ts b/src/main.ts
index ff1100d8..38b9faf7 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -23,11 +23,24 @@ import { PythonEnvironment } from './pythonEnvironment';
 import { DownloadManager } from './models/DownloadManager';
 import { getModelsDirectory } from './utils';
 import { ComfySettings } from './config/comfySettings';
+import dotenv from 'dotenv';
+
+dotenv.config();
 
 let comfyServerProcess: ChildProcess | null = null;
 let isRestarting: boolean = false; // Prevents double restarts TODO(robinhuang): Remove this once we have a better way to handle restarts. https://github.com/Comfy-Org/electron/issues/149
-const host = '127.0.0.1';
-let port = 8188;
+
+/** The host to use for the ComfyUI server. */
+const host = process.env.COMFY_HOST || '127.0.0.1';
+/** The port to use for the ComfyUI server. */
+let port = parseInt(process.env.COMFY_PORT || '-1');
+/**
+ * Whether to use an external server instead of starting one locally.
+ * Only effective if COMFY_PORT is set.
+ * Note: currently used for testing only.
+ */
+const useExternalServer = process.env.USE_EXTERNAL_SERVER === 'true';
+
 let mainWindow: BrowserWindow | null = null;
 let store: Store<StoreType> | null = null;
 const messageQueue: Array<any> = []; // Stores mesaages before renderer is ready.
@@ -188,27 +201,37 @@ if (!gotTheLock) {
       downloadManager = DownloadManager.getInstance(mainWindow!, getModelsDirectory(basePath));
       downloadManager.registerIpcHandlers();
 
-      port = await findAvailablePort(8000, 9999).catch((err) => {
-        log.error(`ERROR: Failed to find available port: ${err}`);
-        throw err;
-      });
-
-      sendProgressUpdate('Setting up Python Environment...');
-      const pythonEnvironment = new PythonEnvironment(pythonInstallPath, appResourcesPath, spawnPythonAsync);
-      await pythonEnvironment.setup();
-      SetupTray(
-        mainWindow,
-        basePath,
-        modelConfigPath,
-        () => {
-          log.info('Resetting install location');
-          fs.rmSync(modelConfigPath);
-          restartApp();
-        },
-        pythonEnvironment
-      );
-      sendProgressUpdate('Starting Comfy Server...');
-      await launchPythonServer(pythonEnvironment.pythonInterpreterPath, appResourcesPath, modelConfigPath, basePath);
+      port =
+        port !== -1
+          ? port
+          : await findAvailablePort(8000, 9999).catch((err) => {
+              log.error(`ERROR: Failed to find available port: ${err}`);
+              throw err;
+            });
+
+      if (!useExternalServer) {
+        sendProgressUpdate('Setting up Python Environment...');
+        const pythonEnvironment = new PythonEnvironment(pythonInstallPath, appResourcesPath, spawnPythonAsync);
+        await pythonEnvironment.setup();
+
+        // TODO: Make tray setup more flexible here as not all actions depend on the python environment.
+        SetupTray(
+          mainWindow,
+          basePath,
+          modelConfigPath,
+          () => {
+            log.info('Resetting install location');
+            fs.rmSync(modelConfigPath);
+            restartApp();
+          },
+          pythonEnvironment
+        );
+        sendProgressUpdate('Starting Comfy Server...');
+        await launchPythonServer(pythonEnvironment.pythonInterpreterPath, appResourcesPath, modelConfigPath, basePath);
+      } else {
+        sendProgressUpdate('Using external server at ' + host + ':' + port);
+        loadComfyIntoMainWindow();
+      }
     } catch (error) {
       log.error(error);
       sendProgressUpdate(COMFY_ERROR_MESSAGE);