Fix GPU preselection (#501)

* Fix frontend can't get detected GPU IPC handler called before registered. * nit * Fix installState never marked as 'started' * Fix test - missing mock
Comfy-Org · Dec 17, 2024 · cdd31f1 · cdd31f1
1 parent e12afde
commit cdd31f1
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/src/handlers/appInfoHandlers.ts b/src/handlers/appInfoHandlers.ts
@@ -1,11 +1,12 @@
 import { app, ipcMain } from 'electron';
 import { IPC_CHANNELS } from '../constants';
+import { useDesktopConfig } from '../store/desktopConfig';
+import type { TorchDeviceType } from '../preload';
+
 /**
- * Handles static information about the app in IPC channels.
+ * Handles information about the app and current state in IPC channels.
  */
 export class AppInfoHandlers {
-  constructor() {}
-
   registerHandlers() {
     ipcMain.handle(IPC_CHANNELS.IS_PACKAGED, () => {
       return app.isPackaged;
@@ -14,5 +15,10 @@ export class AppInfoHandlers {
     ipcMain.handle(IPC_CHANNELS.GET_ELECTRON_VERSION, () => {
       return app.getVersion();
     });
+
+    // Config
+    ipcMain.handle(IPC_CHANNELS.GET_GPU, async (): Promise<TorchDeviceType | undefined> => {
+      return await useDesktopConfig().getAsync('detectedGpu');
+    });
   }
 }
diff --git a/src/main-process/comfyDesktopApp.ts b/src/main-process/comfyDesktopApp.ts
@@ -142,10 +142,6 @@ export class ComfyDesktopApp {
         }
       }
     );
-    // Config
-    ipcMain.handle(IPC_CHANNELS.GET_GPU, async (): Promise<TorchDeviceType | undefined> => {
-      return await useDesktopConfig().getAsync('detectedGpu');
-    });
     // Restart core
     ipcMain.handle(IPC_CHANNELS.RESTART_CORE, async (): Promise<boolean> => {
       if (!this.comfyServer) return false;
@@ -159,8 +155,11 @@ export class ComfyDesktopApp {
    * Install ComfyUI and return the base path.
    */
   static async install(appWindow: AppWindow): Promise<string> {
+    const config = useDesktopConfig();
+    if (!config.get('installState')) config.set('installState', 'started');
+
     const validation = await validateHardware();
-    if (typeof validation?.gpu === 'string') useDesktopConfig().set('detectedGpu', validation.gpu);
+    if (typeof validation?.gpu === 'string') config.set('detectedGpu', validation.gpu);
 
     if (!validation.isValid) {
       await appWindow.loadRenderer('not-supported');

diff --git a/tests/unit/handlers/appinfoHandlers.test.ts b/tests/unit/handlers/appinfoHandlers.test.ts
@@ -4,6 +4,7 @@ import { IPC_CHANNELS } from '../../../src/constants';
 
 jest.mock('electron', () => ({
   ipcMain: {
+    on: jest.fn(),
     handle: jest.fn(),
   },
 }));