Migrate custom nodes

package.json

@@ -12,7 +12,7 @@
   "config": {
     "frontendVersion": "1.6.1",
     "comfyVersion": "0.3.8",
-    "managerCommit": "52b6e4b2eb6481333b29556788f28633e73bd0a9",
+    "managerCommit": "b4e52e65ec87978fc6294e82d55e1b91c5b02d55",
     "uvVersion": "0.5.8"
   },
   "scripts": {
@@ -68,6 +68,7 @@
     "@types/jest": "^29.5.14",
     "@types/node": "^22.10.2",
     "@types/tar": "6.1.13",
+    "@types/tmp": "^0.2.6",
     "@types/wait-on": "^5.3.4",
     "@typescript-eslint/eslint-plugin": "^5.0.0",
     "@typescript-eslint/parser": "^5.0.0",
@@ -107,6 +108,7 @@
     "node-pty": "^1.0.0",
     "systeminformation": "^5.23.5",
     "tar": "^7.4.3",
+    "tmp": "^0.2.3",
     "wait-on": "^8.0.1",
     "yaml": "^2.6.0"
   },

src/constants.ts

@@ -30,6 +30,8 @@ export const IPC_CHANNELS = {
   VALIDATE_COMFYUI_SOURCE: 'validate-comfyui-source',
   SHOW_DIRECTORY_PICKER: 'show-directory-picker',
   INSTALL_COMFYUI: 'install-comfyui',
+  SHOW_TOAST: 'show-toast',
+  LOADED: 'loaded',
   SHOW_CONTEXT_MENU: 'show-context-menu',
   RESTART_CORE: 'restart-core',
   GET_GPU: 'get-gpu',
@@ -92,13 +94,11 @@ export const MigrationItems: MigrationItem[] = [
     label: 'Models',
     description: 'Reference model files from existing ComfyUI installations. (No copy)',
   },
-  // TODO: Decide whether we want to auto-migrate custom nodes, and install their dependencies.
-  // huchenlei: This is a very essential thing for migration experience.
-  // {
-  //   id: 'custom_nodes',
-  //   label: 'Custom Nodes',
-  //   description: 'Reference custom node files from existing ComfyUI installations. (No copy)',
-  // },
+  {
+    id: 'custom_nodes',
+    label: 'Custom Nodes',
+    description: 'Reinstall custom nodes from existing ComfyUI installations.',
+  },
 ] as const;
 
 export const DEFAULT_SERVER_ARGS = {

src/install/installWizard.ts

@@ -25,6 +25,10 @@ export class InstallWizard {
     return !!this.migrationSource && this.migrationItemIds.has('models');
   }
 
+  get shouldMigrateCustomNodes(): boolean {
+    return !!this.migrationSource && this.migrationItemIds.has('custom_nodes');
+  }
+
   get basePath(): string {
     return path.join(this.installOptions.installPath, 'ComfyUI');
   }

src/main-process/appWindow.ts

@@ -97,6 +97,17 @@ export class AppWindow {
     });
   }
 
+  /**
+   * Shows a toast popup in the client
+   */
+  showToast(config: {
+    severity?: 'success' | 'info' | 'warn' | 'error' | 'secondary' | 'contrast' | undefined;
+    summary?: string | undefined;
+    detail?: any | undefined;
+  }): void {
+    this.send(IPC_CHANNELS.SHOW_TOAST, config);
+  }
+
   public onClose(callback: () => void): void {
     this.window.on('close', () => {
       callback();

src/main-process/comfyDesktopApp.ts

@@ -13,12 +13,13 @@ import { InstallOptions, type ElectronContextMenuOptions, type TorchDeviceType }
 import path from 'node:path';
 import { ansiCodes, getModelsDirectory, validateHardware } from '../utils';
 import { DownloadManager } from '../models/DownloadManager';
-import { VirtualEnvironment } from '../virtualEnvironment';
+import { ProcessCallbacks, VirtualEnvironment } from '../virtualEnvironment';
 import { InstallWizard } from '../install/installWizard';
 import { Terminal } from '../shell/terminal';
-import { useDesktopConfig } from '../store/desktopConfig';
+import { DesktopConfig, useDesktopConfig } from '../store/desktopConfig';
 import { InstallationValidator } from '../install/installationValidator';
 import { restoreCustomNodes } from '../services/backup';
+import { CmCli } from '../services/cmCli';
 
 export class ComfyDesktopApp {
   public comfyServer: ComfyServer | null = null;
@@ -183,6 +184,9 @@ export class ComfyDesktopApp {
           .then(() => {
             useDesktopConfig().set('installState', 'installed');
             appWindow.maximize();
+            if (installWizard.shouldMigrateCustomNodes && installWizard.migrationSource) {
+              useDesktopConfig().set('migrateCustomNodesFrom', installWizard.migrationSource);
+            }
             resolve(installWizard.basePath);
           })
           .catch(reject);
@@ -213,7 +217,7 @@ export class ComfyDesktopApp {
     const selectedDevice = config.get('selectedDevice');
     const virtualEnvironment = new VirtualEnvironment(this.basePath, selectedDevice);
 
-    await virtualEnvironment.create({
+    const processCallbacks: ProcessCallbacks = {
       onStdout: (data) => {
         log.info(data.replaceAll(ansiCodes, ''));
         this.appWindow.send(IPC_CHANNELS.LOG_MESSAGE, data);
@@ -222,7 +226,11 @@ export class ComfyDesktopApp {
         log.error(data.replaceAll(ansiCodes, ''));
         this.appWindow.send(IPC_CHANNELS.LOG_MESSAGE, data);
       },
-    });
+    };
+
+    await virtualEnvironment.create(processCallbacks);
+
+    const customNodeMigrationError = await this.migrateCustomNodes(config, virtualEnvironment, processCallbacks);
 
     if (!config.get('Comfy-Desktop.RestoredCustomNodes', false)) {
       try {
@@ -238,6 +246,36 @@ export class ComfyDesktopApp {
     this.comfyServer = new ComfyServer(this.basePath, serverArgs, virtualEnvironment, this.appWindow);
     await this.comfyServer.start();
     this.initializeTerminal(virtualEnvironment);
+
+    return () => {
+      if (customNodeMigrationError) {
+        this.appWindow.showToast({
+          summary: 'Failed to migrate custom nodes',
+          detail: customNodeMigrationError,
+          severity: 'error',
+        });
+      }
+    };
+  }
+
+  async migrateCustomNodes(config: DesktopConfig, virtualEnvironment: VirtualEnvironment, callbacks: ProcessCallbacks) {
+    const customNodeMigrationPath = config.get('migrateCustomNodesFrom');
+    let customNodeMigrationError: string | null = null;
+    if (customNodeMigrationPath) {
+      log.info('Migrating custom nodes from: ', customNodeMigrationPath);
+      try {
+        const cmCli = new CmCli(virtualEnvironment);
+        await cmCli.restoreCustomNodes(customNodeMigrationPath, callbacks);
+      } catch (error) {
+        log.error('Error migrating custom nodes', error);
+        customNodeMigrationError =
+          error instanceof Error ? error.message : typeof error === 'string' ? error : 'Error migrating custom nodes.';
+      } finally {
+        // Always remove the flag so the user doesnt get stuck here
+        config.delete('migrateCustomNodesFrom');
+      }
+    }
+    return customNodeMigrationError;
   }
 
   static async create(appWindow: AppWindow): Promise<ComfyDesktopApp> {

src/main.ts

@@ -94,11 +94,17 @@ async function startApp() {
       delete extraServerArgs.listen;
       delete extraServerArgs.port;
 
+      let onReady: (() => void) | null = null;
       if (!useExternalServer) {
-        await comfyDesktopApp.startComfyServer({ host, port, extraServerArgs });
+        onReady = await comfyDesktopApp.startComfyServer({ host, port, extraServerArgs });
       }
       appWindow.sendServerStartProgress(ProgressStatus.READY);
+
+      const waitLoad = Promise.withResolvers();
+      ipcMain.handle(IPC_CHANNELS.LOADED, waitLoad.resolve);
       await appWindow.loadComfyUI({ host, port, extraServerArgs });
+      await waitLoad.promise;
+      onReady?.();
     } catch (error) {
       log.error('Unhandled exception during app startup', error);
       appWindow.sendServerStartProgress(ProgressStatus.ERROR);

src/preload.ts

@@ -223,6 +223,20 @@ const electronAPI = {
   installComfyUI: (installOptions: InstallOptions) => {
     ipcRenderer.send(IPC_CHANNELS.INSTALL_COMFYUI, installOptions);
   },
+  /**
+   * Shows a toast popup with the supplied config
+   */
+  onShowToast: (callback: (config: { message: string; type: string }) => void) => {
+    ipcRenderer.on(IPC_CHANNELS.SHOW_TOAST, (_event, value) => {
+      callback(value);
+    });
+  },
+  /**
+   * Trigged by the frontend app when it is loaded
+   */
+  loaded: (): Promise<string> => {
+    return ipcRenderer.invoke(IPC_CHANNELS.LOADED);
+  },
   /**
    * Opens native context menus.
    *

src/services/cmCli.ts

@@ -0,0 +1,92 @@
+import log from 'electron-log/main';
+import path from 'path';
+import { getAppResourcesPath } from '../install/resourcePaths';
+import { ProcessCallbacks, VirtualEnvironment } from '../virtualEnvironment';
+import { fileSync } from 'tmp';
+
+export class CmCli {
+  private cliPath: string;
+  private virtualEnvironment: VirtualEnvironment;
+
+  constructor(virtualEnvironment: VirtualEnvironment) {
+    this.virtualEnvironment = virtualEnvironment;
+    this.cliPath = path.join(getAppResourcesPath(), 'ComfyUI', 'custom_nodes', 'ComfyUI-Manager', 'cm-cli.py');
+  }
+
+  private async _runCommandAsync(
+    args: string[],
+    callbacks?: ProcessCallbacks,
+    env?: Record<string, string>,
+    cwd?: string
+  ): Promise<{ exitCode: number | null }> {
+    const cmd = [this.cliPath, ...args];
+    return await this.virtualEnvironment.runPythonCommandAsync(cmd, callbacks, env, cwd);
+  }
+
+  public async runCommandAsync(
+    args: string[],
+    callbacks?: ProcessCallbacks,
+    env: Record<string, string> = {},
+    checkExit: boolean = true,
+    cwd?: string
+  ) {
+    let output = '';
+    let error = '';
+    const { exitCode } = await this._runCommandAsync(
+      args,
+      {
+        onStdout: (message) => {
+          output += message;
+          callbacks?.onStdout?.(message);
+        },
+        onStderr: (message) => {
+          console.warn('[warn]', message);
+          error += message;
+          callbacks?.onStderr?.(message);
+        },
+      },
+      {
+        COMFYUI_PATH: this.virtualEnvironment.venvRootPath,
+        ...env,
+      },
+      cwd
+    );
+
+    if (checkExit && exitCode !== 0) {
+      throw new Error(`Error calling cm-cli: \nExit code: ${exitCode}\nOutput:${output}\n\nError:${error}`);
+    }
+
+    return output;
+  }
+
+  public async restoreCustomNodes(fromComfyDir: string, callbacks: ProcessCallbacks) {
+    const tmpFile = fileSync({ postfix: '.json' });
+    try {
+      log.debug('Using temp file: ' + tmpFile.name);
+      await this.saveSnapshot(fromComfyDir, tmpFile.name, callbacks);
+      await this.restoreSnapshot(tmpFile.name, callbacks);
+    } finally {
+      tmpFile?.removeCallback();
+    }
+  }
+
+  public async saveSnapshot(fromComfyDir: string, outFile: string, callbacks: ProcessCallbacks): Promise<void> {
+    const output = await this.runCommandAsync(
+      ['save-snapshot', '--output', outFile, '--no-full-snapshot'],
+      callbacks,
+      {
+        COMFYUI_PATH: fromComfyDir,
+        PYTHONPATH: fromComfyDir,
+      },
+      true,
+      fromComfyDir
+    );
+    log.info(output);
+  }
+
+  public async restoreSnapshot(snapshotFile: string, callbacks: ProcessCallbacks) {
+    log.info('Restoring snapshot ' + snapshotFile);
+    const output = await this.runCommandAsync(['restore-snapshot', snapshotFile], callbacks);
+    log.info(output);
+  }
+}

src/store/desktopConfig.ts

@@ -32,6 +32,11 @@ export class DesktopConfig {
     return value === undefined ? this.#store.delete(key) : this.#store.set(key, value);
   }
 
+  /** @inheritdoc {@link ElectronStore.delete} */
+  delete<Key extends keyof DesktopSettings>(key: Key) {
+    this.#store.delete(key);
+  }
+
   /**
    * Static factory method. Loads the config from disk.
    * @param shell Shell environment that can open file and folder views for the user

src/store/index.ts

@@ -18,6 +18,10 @@ export type DesktopSettings = {
    * in the yaml config.
    */
   installState?: 'started' | 'installed' | 'upgraded';
+  /**
+   * The path to the migration installation to migrate custom nodes from
+   */
+  migrateCustomNodesFrom?: string;
   /**
    * The last GPU that was detected during hardware validation.
    * Allows manual override of some install behaviour.

src/virtualEnvironment.ts

@@ -8,7 +8,7 @@ import os from 'node:os';
 import { getDefaultShell } from './shell/util';
 import type { TorchDeviceType } from './preload';
 
-type ProcessCallbacks = {
+export type ProcessCallbacks = {
   onStdout?: (data: string) => void;
   onStderr?: (data: string) => void;
 };
@@ -210,15 +210,19 @@ export class VirtualEnvironment {
    */
   public async runPythonCommandAsync(
     args: string[],
-    callbacks?: ProcessCallbacks
+    callbacks?: ProcessCallbacks,
+    env?: Record<string, string>,
+    cwd?: string
   ): Promise<{ exitCode: number | null }> {
     return this.runCommandAsync(
       this.pythonInterpreterPath,
       args,
       {
+        ...env,
         PYTHONIOENCODING: 'utf8',
       },
-      callbacks
+      callbacks,
+      cwd
     );
   }
 
@@ -276,11 +280,12 @@ export class VirtualEnvironment {
     command: string,
     args: string[],
     env: Record<string, string>,
-    callbacks?: ProcessCallbacks
+    callbacks?: ProcessCallbacks,
+    cwd?: string
   ): ChildProcess {
     log.info(`Running command: ${command} ${args.join(' ')} in ${this.venvRootPath}`);
     const childProcess: ChildProcess = spawn(command, args, {
-      cwd: this.venvRootPath,
+      cwd: cwd ?? this.venvRootPath,
       env: {
         ...process.env,
         ...env,
@@ -306,10 +311,11 @@ export class VirtualEnvironment {
     command: string,
     args: string[],
     env: Record<string, string>,
-    callbacks?: ProcessCallbacks
+    callbacks?: ProcessCallbacks,
+    cwd?: string
   ): Promise<{ exitCode: number | null }> {
     return new Promise((resolve, reject) => {
-      const childProcess = this.runCommand(command, args, env, callbacks);
+      const childProcess = this.runCommand(command, args, env, callbacks, cwd);
 
       childProcess.on('close', (code) => {
         resolve({ exitCode: code });