Add progress report bar. (#20)

* Add progress report. * Refactor. * Add comment. * Update.
Comfy-Org · Sep 13, 2024 · 0a374cf · 0a374cf
1 parent e052778
commit 0a374cf
Show file tree

Hide file tree

Showing 7 changed files with 347 additions and 278 deletions.
diff --git a/index.html b/index.html
@@ -3,11 +3,41 @@
   <head>
     <meta charset="UTF-8" />
     <title>ComfyUI</title>
-
+    <style>
+      body {
+        font-family: Arial, sans-serif;
+        display: flex;
+        flex-direction: column;
+        justify-content: center;
+        align-items: center;
+        height: 100vh;
+        margin: 0;
+        background-color: #f0f0f0;
+      }
+      #loading-text {
+        margin-bottom: 20px;
+        text-align: center;
+      }
+      #progress-bar {
+        width: 300px;
+        height: 20px;
+        background-color: #e0e0e0;
+        border-radius: 10px;
+        overflow: hidden;
+      }
+      #progress {
+        width: 0%;
+        height: 100%;
+        background-color: #09f;
+        transition: width 0.5s ease;
+      }
+    </style>
   </head>
   <body>
-    <h1>Python Server is Loading</h1>
-    <p>⏳</p>
+    <div id="loading-text">Initializing...</div>
+    <div id="progress-bar">
+      <div id="progress"></div>
+    </div>
     <script type="module" src="/src/renderer.ts"></script>
   </body>
 </html>
diff --git a/package.json b/package.json
@@ -16,6 +16,7 @@
     "make:assets:amd": "cd assets && comfy-cli --skip-prompt --here install --fast-deps --amd --manager-url https://github.com/Comfy-Org/manager-core && comfy-cli --here standalone",
     "make:assets:cpu": "cd assets && comfy-cli --skip-prompt --here install --fast-deps --cpu --manager-url https://github.com/Comfy-Org/manager-core && comfy-cli --here standalone && node ../scripts/env.mjs",
     "make:assets:nvidia": "cd assets && comfy-cli --skip-prompt --here install --fast-deps --nvidia --manager-url https://github.com/Comfy-Org/manager-core && comfy-cli --here standalone",
+    "make:assets:macos": "cd assets && comfy-cli --skip-prompt --here install --fast-deps --nvidia --manager-url https://github.com/Comfy-Org/manager-core && comfy-cli --here standalone",
     "notarize": "node debug/notarize.js",
     "package": "electron-forge package",
     "publish": "electron-forge publish",

diff --git a/src/constants.ts b/src/constants.ts
@@ -0,0 +1,5 @@
+export const IPC_CHANNELS = {
+    LOADING_PROGRESS: 'loading-progress',
+  };
+
+  export const ELECTRON_BRIDGE_API = 'electronAPI';
diff --git a/src/main.ts b/src/main.ts
@@ -4,45 +4,43 @@ import fs from 'fs'
 import net from 'node:net';
 import path from 'node:path';
 import { SetupTray } from './tray';
-
+import { IPC_CHANNELS } from './constants';
 import dotenv from "dotenv";
-import { app, BrowserWindow, webContents } from 'electron';
+import { app, BrowserWindow, webContents, ipcMain, screen } from 'electron';
 // Handle creating/removing shortcuts on Windows when installing/uninstalling.
 import('electron-squirrel-startup').then(ess => {
   const {default: check} = ess;
   if (check) {
     app.quit();
   }
 });
-import tar from 'tar';
-import { create } from 'node:domain';
 
 let pythonProcess: ChildProcess | null = null;
 const host = '127.0.0.1'; // Replace with the desired IP address
 const port = 8188; // Replace with the port number your server is running on
+let mainWindow: BrowserWindow | null;
+
 
-const createWindow = () => {
-  // Create the browser window.
-  const mainWindow = new BrowserWindow({
+const createWindow = async () => {
+    const primaryDisplay = screen.getPrimaryDisplay();
+    const { width, height } = primaryDisplay.workAreaSize;
+  mainWindow = new BrowserWindow({
     title: 'ComfyUI',
-    width: 800,
-    height: 600,
+    width: width,
+    height: height,
     webPreferences: {
       preload: path.join(__dirname, 'preload.js'),
       nodeIntegration: true, // Enable Node.js integration
-      contextIsolation: false,
+      contextIsolation: true,
     },
-
   });
-
-  // Load the UI from the Python server's URL
-  //mainWindow.loadURL('http://localhost:8188/');
-
+
   // and load the index.html of the app.
   if (MAIN_WINDOW_VITE_DEV_SERVER_URL) {
-    mainWindow.loadURL(MAIN_WINDOW_VITE_DEV_SERVER_URL);
+    console.log('Loading Vite Dev Server');
+    await mainWindow.loadURL(MAIN_WINDOW_VITE_DEV_SERVER_URL);
   } else {
-    mainWindow.loadFile(path.join(__dirname, `../renderer/${MAIN_WINDOW_VITE_NAME}/index.html`));
+    await mainWindow.loadFile(path.join(__dirname, `../renderer/${MAIN_WINDOW_VITE_NAME}/index.html`));
   }
 
   // Set up the System Tray Icon for all platforms
@@ -99,8 +97,8 @@ const isPortInUse = (host: string, port: number): Promise<boolean> => {
 
 // Launch Python Server Variables
 const maxFailWait: number = 50 * 1000; // 50seconds
-let currentWaitTime: number = 0;
-let spawnServerTimeout: NodeJS.Timeout = null;
+let currentWaitTime = 0;
+const spawnServerTimeout: NodeJS.Timeout = null;
 
 const launchPythonServer = async (args: {userResourcesPath: string, appResourcesPath: string}) => {
   const {userResourcesPath, appResourcesPath} = args;
@@ -155,12 +153,12 @@ const launchPythonServer = async (args: {userResourcesPath: string, appResources
     } catch {
       console.log('Running one-time python installation on first startup...');
       // clean up any possible existing non-functional python env
-      try {
-        await fsPromises.rm(pythonRootPath, {recursive: true});
-      } catch {null;}
+    //   try {
+    //     await fsPromises.rm(pythonRootPath, {recursive: true});
+    //   } catch {null;}
 
-      const pythonTarPath = path.join(appResourcesPath, 'python.tgz');
-      await tar.extract({file: pythonTarPath, cwd: userResourcesPath, strict: true});
+    //   const pythonTarPath = path.join(appResourcesPath, 'python.tgz');
+    //   await tar.extract({file: pythonTarPath, cwd: userResourcesPath, strict: true});
 
       const wheelsPath = path.join(pythonRootPath, 'wheels');
       // TODO: report space bug to uv upstream, then revert below mac fix
@@ -194,12 +192,13 @@ const launchPythonServer = async (args: {userResourcesPath: string, appResources
       }
       const isReady = await isPortInUse(host, port);
       if (isReady) {
+        sendProgressUpdate(90, 'Finishing...');    
         console.log('Python server is ready');
         // Start the Heartbeat listener, send connected message to Renderer and resolve promise.
         serverHeartBeatReference = setInterval(serverHeartBeat, serverHeartBeatInterval);
         webContents.getAllWebContents()[0].send("python-server-status", "active");
         //For now just replace the source of the main window to the python server
-        webContents.getAllWebContents()[0].loadURL('http://localhost:8188/');
+        setTimeout( () => webContents.getAllWebContents()[0].loadURL('http://localhost:8188/'), 1000);
         clearTimeout(spawnServerTimeout);
         resolve();
       } else {
@@ -225,7 +224,6 @@ app.on('ready', async () => {
     userResourcesPath: path.join(app.getAppPath(), 'assets'),
     appResourcesPath: path.join(app.getAppPath(), 'assets'),
   }
-
   console.log(`userResourcesPath: ${userResourcesPath}`);
   console.log(`appResourcesPath: ${appResourcesPath}`);
 
@@ -241,17 +239,27 @@ app.on('ready', async () => {
     // if user-specific resources dir already exists, that is fine
   }
   try {
-    createWindow();
+    await createWindow();
+    sendProgressUpdate(20, 'Setting up comfy environment...');
     createComfyDirectories();
+    setTimeout(() => sendProgressUpdate(40, 'Starting Comfy Server...'), 1000);
     await launchPythonServer({userResourcesPath, appResourcesPath});
   } catch (error) {
     console.error(error);
+    sendProgressUpdate(0, 'Failed to start Comfy Server');
   }
 });
 
+function sendProgressUpdate(percentage: number, status: string) {
+    if (mainWindow) {
+        console.log('Sending progress update to renderer ' + status);
+        mainWindow.webContents.send(IPC_CHANNELS.LOADING_PROGRESS, { percentage, status });
+    }
+  }
+
 const killPythonServer = () => {
   console.log('Python server:', pythonProcess);
-  return new Promise<void>(async(resolve, reject) => {
+  return new Promise<void>((resolve, reject) => {
     if (pythonProcess) {
       try {
         const result:boolean = pythonProcess.kill(); //false if kill did not succeed sucessfully
@@ -272,6 +280,7 @@ const killPythonServer = () => {
 
 type DirectoryStructure = (string | [string, string[]])[];
 
+// Create directories needed by ComfyUI in the user's data directory.
 function createComfyDirectories(): void {
   const userDataPath: string = app.getPath('userData');
   const directories: DirectoryStructure = [
@@ -357,6 +366,3 @@ app.on('activate', () => {
     createWindow();
   }
 });
-
-// In this file you can include the rest of your app's specific main process
-// code. You can also put them in separate files and import them here.
diff --git a/src/preload.ts b/src/preload.ts
@@ -1,2 +1,15 @@
-// See the Electron documentation for details on how to use preload scripts:
-// https://www.electronjs.org/docs/latest/tutorial/process-model#preload-scripts
+// preload.ts
+
+import { contextBridge, ipcRenderer } from 'electron';
+import { IPC_CHANNELS, ELECTRON_BRIDGE_API } from './constants';
+
+const electronAPI = {
+  onProgressUpdate: (callback: (update: { percentage: number; status: string }) => void) => {
+    ipcRenderer.on(IPC_CHANNELS.LOADING_PROGRESS, (_event, value) => {
+      console.log(`Received ${IPC_CHANNELS.LOADING_PROGRESS} event`, value);
+      callback(value);
+    });
+  },
+};
+
+contextBridge.exposeInMainWorld(ELECTRON_BRIDGE_API, electronAPI);
diff --git a/src/renderer.ts b/src/renderer.ts
@@ -27,5 +27,33 @@
  */
 
 import './index.css';
-
+import { IPC_CHANNELS, ELECTRON_BRIDGE_API } from './constants';
 console.log('👋 This message is being logged by "renderer.ts", included via Vite');
+
+interface ProgressUpdate {
+  percentage: number;
+  status: string;
+}
+
+const progressBar = document.getElementById('progress') as HTMLElement;
+const loadingText = document.getElementById('loading-text') as HTMLElement;
+
+function updateProgress({ percentage, status }: ProgressUpdate) {
+  console.log(`Updating progress: ${percentage}%, ${status}`);
+  progressBar.style.width = `${percentage}%`;
+  loadingText.textContent = status;
+
+  if (percentage === 100) {
+    loadingText.textContent = 'ComfyUI is ready!';
+  }
+}
+
+if (ELECTRON_BRIDGE_API in window) {
+  console.log(`${ELECTRON_BRIDGE_API} found, setting up listeners`);
+  (window as any).electronAPI.onProgressUpdate((update: ProgressUpdate) => {
+    console.log("Received loading progress", update);
+    updateProgress(update);
+  });
+} else {
+  console.error(`${ELECTRON_BRIDGE_API} not found in window object`);
+}