Skip to content

Commit

Permalink
Migrate custom nodes (#474)
Browse files Browse the repository at this point in the history
* Add custom node migration

* Merge fixes

* Update manager commit

* Remove error toastr
  • Loading branch information
pythongosssss authored Dec 23, 2024
1 parent 03d12d4 commit e3d6fa2
Show file tree
Hide file tree
Showing 9 changed files with 176 additions and 21 deletions.
4 changes: 3 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
"config": {
"frontendVersion": "1.6.1",
"comfyVersion": "0.3.9",
"managerCommit": "52b6e4b2eb6481333b29556788f28633e73bd0a9",
"managerCommit": "b4e52e65ec87978fc6294e82d55e1b91c5b02d55",
"uvVersion": "0.5.8"
},
"scripts": {
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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"
},
Expand Down
12 changes: 5 additions & 7 deletions src/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -92,13 +92,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 = {
Expand Down
4 changes: 4 additions & 0 deletions src/install/installWizard.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
}
Expand Down
45 changes: 40 additions & 5 deletions src/main-process/comfyDesktopApp.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { app, dialog, ipcMain } from 'electron';
import { app, dialog, ipcMain, Notification } from 'electron';
import log from 'electron-log/main';
import * as Sentry from '@sentry/electron/main';
import { graphics } from 'systeminformation';
Expand All @@ -13,12 +13,13 @@ import { InstallOptions, type ElectronContextMenuOptions } from '../preload';
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;
Expand Down Expand Up @@ -186,6 +187,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);
Expand Down Expand Up @@ -216,7 +220,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);
Expand All @@ -225,7 +229,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 {
Expand All @@ -241,6 +249,33 @@ export class ComfyDesktopApp {
this.comfyServer = new ComfyServer(this.basePath, serverArgs, virtualEnvironment, this.appWindow);
await this.comfyServer.start();
this.initializeTerminal(virtualEnvironment);

if (customNodeMigrationError) {
new Notification({
title: 'Failed to migrate custom nodes',
body: customNodeMigrationError,
}).show();
}
}

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> {
Expand Down
92 changes: 92 additions & 0 deletions src/services/cmCli.ts
Original file line number Diff line number Diff line change
@@ -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);
}
}
5 changes: 5 additions & 0 deletions src/store/desktopConfig.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions src/store/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
20 changes: 13 additions & 7 deletions src/virtualEnvironment.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import os, { EOL } 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;
};
Expand Down Expand Up @@ -215,15 +215,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
);
}

Expand Down Expand Up @@ -281,11 +285,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,
Expand All @@ -311,10 +316,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 });
Expand Down
11 changes: 10 additions & 1 deletion yarn.lock
Original file line number Diff line number Diff line change
Expand Up @@ -555,6 +555,7 @@ __metadata:
"@types/jest": "npm:^29.5.14"
"@types/node": "npm:^22.10.2"
"@types/tar": "npm:6.1.13"
"@types/tmp": "npm:^0.2.6"
"@types/wait-on": "npm:^5.3.4"
"@typescript-eslint/eslint-plugin": "npm:^5.0.0"
"@typescript-eslint/parser": "npm:^5.0.0"
Expand All @@ -577,6 +578,7 @@ __metadata:
rimraf: "npm:^6.0.1"
systeminformation: "npm:^5.23.5"
tar: "npm:^7.4.3"
tmp: "npm:^0.2.3"
ts-jest: "npm:^29.2.5"
ts-node: "npm:^10.0.0"
typescript: "npm:^5.7.2"
Expand Down Expand Up @@ -3741,6 +3743,13 @@ __metadata:
languageName: node
linkType: hard

"@types/tmp@npm:^0.2.6":
version: 0.2.6
resolution: "@types/tmp@npm:0.2.6"
checksum: 10c0/a11bfa2cd8eaa6c5d62f62a3569192d7a2c28efdc5c17af0b0551db85816b2afc8156f3ca15ac76f0b142ae1403f04f44279871424233a1f3390b2e5fc828cd0
languageName: node
linkType: hard

"@types/verror@npm:^1.10.3":
version: 1.10.10
resolution: "@types/verror@npm:1.10.10"
Expand Down Expand Up @@ -12920,7 +12929,7 @@ __metadata:
languageName: node
linkType: hard

"tmp@npm:^0.2.0":
"tmp@npm:^0.2.0, tmp@npm:^0.2.3":
version: 0.2.3
resolution: "tmp@npm:0.2.3"
checksum: 10c0/3e809d9c2f46817475b452725c2aaa5d11985cf18d32a7a970ff25b568438e2c076c2e8609224feef3b7923fa9749b74428e3e634f6b8e520c534eef2fd24125
Expand Down

0 comments on commit e3d6fa2

Please sign in to comment.