Skip to content

Commit

Permalink
Refactor install.
Browse files Browse the repository at this point in the history
  • Loading branch information
robinjhuang committed Nov 12, 2024
1 parent b65ea08 commit fca7b71
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 42 deletions.
3 changes: 3 additions & 0 deletions src/config/comfyConfigManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@ import log from 'electron-log/main';

export type DirectoryStructure = (string | DirectoryStructure)[];

/**
* Responsible for creating the ComfyUI directory structure.
*/
export class ComfyConfigManager {
private static readonly DEFAULT_DIRECTORIES: DirectoryStructure = [
'custom_nodes',
Expand Down
4 changes: 4 additions & 0 deletions src/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@ export const IPC_CHANNELS = {
OPEN_DIALOG: 'open-dialog',
FIRST_TIME_SETUP_COMPLETE: 'first-time-setup-complete',
DEFAULT_INSTALL_LOCATION: 'default-install-location',
// Installation state
INSTALLATION_STATE_CHANGED: 'installation-state-changed',
GET_INSTALLATION_STATE: 'get-installation-state',
// End Installation state
DOWNLOAD_PROGRESS: 'download-progress',
START_DOWNLOAD: 'start-download',
PAUSE_DOWNLOAD: 'pause-download',
Expand Down
109 changes: 109 additions & 0 deletions src/install/install.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
import { app, ipcMain } from 'electron';
import log from 'electron-log/main';
import { IPC_CHANNELS } from '../constants';
import { createModelConfigFiles, getModelConfigPath } from '../config/extra_model_config';
import fs from 'fs';
import { AppWindow } from '../main-process/appWindow';
import { ComfyConfigManager } from '../config/comfyConfigManager';

type InstallationState =
| { status: 'NOT_INSTALLED' }
| { status: 'ASKING_FOR_DIRECTORY'; defaultLocation: string; validationErrorMessage?: string }
| {
status: 'INSTALLING';
location: string;
}
| {
status: 'READY';
location: string; // Must exist!
}
| {
status: 'ERROR';
error: Error; // Must exist!
};

/**
* This class is responsible for handling the installation of ComfyUI.
* It registers an IPC handler for the current installation state.
* It also sends any changes to the state to the renderer process.
*/
export class ComfyUIInstall {
private static instance: ComfyUIInstall;
private installLocation: string;
private defaultInstallLocation: string;
private state: InstallationState = { status: 'NOT_INSTALLED' };
private appWindow: AppWindow;

private constructor(appWindow: AppWindow) {
const defaultInstallLocation = app.getPath('documents');
ipcMain.handle(IPC_CHANNELS.DEFAULT_INSTALL_LOCATION, () => defaultInstallLocation); // TODO: Remove after migration.
ipcMain.handle(IPC_CHANNELS.GET_INSTALLATION_STATE, () => this.state);
this.installLocation = defaultInstallLocation;
this.defaultInstallLocation = defaultInstallLocation;
this.appWindow = appWindow;
}

static get(appWindow: AppWindow): ComfyUIInstall {
if (!ComfyUIInstall.instance) {
ComfyUIInstall.instance = new ComfyUIInstall(appWindow);
}
return ComfyUIInstall.instance;
}

setState(state: InstallationState): void {
log.info('Setting installation state:', state);
this.state = state;
this.appWindow.send(IPC_CHANNELS.INSTALLATION_STATE_CHANGED, state);
}

public async install(): Promise<void> {
const firstTimeSetup = await this.isFirstTimeSetup();
log.info('First time setup:', firstTimeSetup);
if (!firstTimeSetup) {
this.appWindow.send(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null); // TODO:Remove after migration.
return;
}
this.setState({ status: 'ASKING_FOR_DIRECTORY', defaultLocation: this.defaultInstallLocation });
this.appWindow.send(IPC_CHANNELS.SHOW_SELECT_DIRECTORY, null); // TODO:Remove after migration.
while (this.state.status === 'ASKING_FOR_DIRECTORY') {
const selectedDirectory = await this.selectInstallDirectory();
const { valid, errorMessage } = await this.isValidComfyDirectory(selectedDirectory);
if (!valid) {
this.setState({ ...this.state, validationErrorMessage: errorMessage });
return;
} else {
this.setState({ status: 'INSTALLING', location: selectedDirectory });
this.installLocation = selectedDirectory;
}
}

const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(this.installLocation);
const modelConfigPath = getModelConfigPath();
await createModelConfigFiles(modelConfigPath, actualComfyDirectory);

this.setState({ status: 'READY', location: this.installLocation });
}

/**
* Check if the user has completed the first time setup wizard.
* This means the extra_models_config.yaml file exists in the user's data directory.
*/
private async isFirstTimeSetup(): Promise<boolean> {
const extraModelsConfigPath = getModelConfigPath();
return !fs.existsSync(extraModelsConfigPath);
}

private async selectInstallDirectory(): Promise<string> {
return new Promise((resolve, reject) => {
ipcMain.on(IPC_CHANNELS.SELECTED_DIRECTORY, (_event, value: string) => {
log.info('Directory selected:', value);
resolve(value);
});
});
}

async isValidComfyDirectory(selectedDirectory: string): Promise<{ valid: boolean; errorMessage?: string }> {
//TODO: Implement
return { valid: true };
}
}
47 changes: 5 additions & 42 deletions src/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,24 @@ import axios from 'axios';
import path from 'node:path';
import { SetupTray } from './tray';
import { IPC_CHANNELS, SENTRY_URL_ENDPOINT, ProgressStatus } from './constants';
import { app, BrowserWindow, dialog, ipcMain, shell } from 'electron';
import { app, dialog, ipcMain } from 'electron';
import log from 'electron-log/main';
import * as Sentry from '@sentry/electron/main';
import * as net from 'net';
import { graphics } from 'systeminformation';
import { createModelConfigFiles, getModelConfigPath, readBasePathFromConfig } from './config/extra_model_config';
import { getModelConfigPath } from './config/extra_model_config';
import todesktop from '@todesktop/runtime';
import { PythonEnvironment } from './pythonEnvironment';
import { DownloadManager } from './models/DownloadManager';
import { getModelsDirectory } from './utils';
import { ComfySettings } from './config/comfySettings';
import dotenv from 'dotenv';
import { buildMenu } from './menu/menu';
import { ComfyConfigManager } from './config/comfyConfigManager';
import { AppWindow } from './main-process/appWindow';
import { getAppResourcesPath, getBasePath, getPythonInstallPath } from './install/resourcePaths';
import { PathHandlers } from './handlers/pathHandlers';
import { AppInfoHandlers } from './handlers/appInfoHandlers';
import { ComfyUIInstall } from './install/install';

dotenv.config();

Expand Down Expand Up @@ -157,13 +157,8 @@ if (!gotTheLock) {
});
});

ipcMain.on(IPC_CHANNELS.OPEN_DEV_TOOLS, () => {
appWindow.openDevTools();
});
ipcMain.handle(IPC_CHANNELS.IS_FIRST_TIME_SETUP, () => {
return isFirstTimeSetup();
});
await handleFirstTimeSetup();
const install = ComfyUIInstall.get(appWindow);
install.install();
const basePath = await getBasePath();
const pythonInstallPath = await getPythonInstallPath();
if (!basePath || !pythonInstallPath) {
Expand Down Expand Up @@ -551,38 +546,6 @@ function findAvailablePort(startPort: number, endPort: number): Promise<number>
tryPort(startPort);
});
}
/**
* Check if the user has completed the first time setup wizard.
* This means the extra_models_config.yaml file exists in the user's data directory.
*/
function isFirstTimeSetup(): boolean {
const extraModelsConfigPath = getModelConfigPath();
return !fs.existsSync(extraModelsConfigPath);
}

async function selectedInstallDirectory(): Promise<string> {
return new Promise((resolve, reject) => {
ipcMain.on(IPC_CHANNELS.SELECTED_DIRECTORY, (_event, value) => {
log.info('User selected to install ComfyUI in:', value);
resolve(value);
});
});
}

async function handleFirstTimeSetup() {
const firstTimeSetup = isFirstTimeSetup();
log.info('First time setup:', firstTimeSetup);
if (firstTimeSetup) {
appWindow.send(IPC_CHANNELS.SHOW_SELECT_DIRECTORY, null);
const selectedDirectory = await selectedInstallDirectory();
const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(selectedDirectory);

const modelConfigPath = getModelConfigPath();
await createModelConfigFiles(modelConfigPath, actualComfyDirectory);
} else {
appWindow.send(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null);
}
}

/**
* Rotate old log files by adding a timestamp to the end of the file.
Expand Down

0 comments on commit fca7b71

Please sign in to comment.