From 7658a4e85d618594f4d41a10c1b63d0ae12f4cbd Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Tue, 5 Nov 2024 15:45:04 -0800 Subject: [PATCH] Refactor creating ComfyUI directories + unit tests. (#185) --- src/__tests__/unit/comfyConfigManager.test.ts | 165 +++++++++++++++++ src/config/comfyConfigManager.ts | 142 ++++++++++++++ src/main.ts | 174 +----------------- src/menu/menu.ts | 32 ++++ 4 files changed, 346 insertions(+), 167 deletions(-) create mode 100644 src/__tests__/unit/comfyConfigManager.test.ts create mode 100644 src/config/comfyConfigManager.ts create mode 100644 src/menu/menu.ts diff --git a/src/__tests__/unit/comfyConfigManager.test.ts b/src/__tests__/unit/comfyConfigManager.test.ts new file mode 100644 index 00000000..c8e90707 --- /dev/null +++ b/src/__tests__/unit/comfyConfigManager.test.ts @@ -0,0 +1,165 @@ +import fs from 'fs'; +import { ComfyConfigManager, DirectoryStructure } from '../../config/comfyConfigManager'; + +// Mock the fs module +jest.mock('fs'); +jest.mock('electron-log/main', () => ({ + info: jest.fn(), + error: jest.fn(), + warn: jest.fn(), +})); + +describe('ComfyConfigManager', () => { + // Reset all mocks before each test + beforeEach(() => { + jest.clearAllMocks(); + (fs.existsSync as jest.Mock).mockReset(); + (fs.mkdirSync as jest.Mock).mockReset(); + (fs.writeFileSync as jest.Mock).mockReset(); + (fs.renameSync as jest.Mock).mockReset(); + }); + + describe('setUpComfyUI', () => { + it('should use existing directory when it contains ComfyUI structure', () => { + // Mock isComfyUIDirectory to return true for the input path + (fs.existsSync as jest.Mock).mockImplementation((path: string) => { + const requiredDirs = [ + '/existing/ComfyUI/models', + '/existing/ComfyUI/input', + '/existing/ComfyUI/user', + '/existing/ComfyUI/output', + '/existing/ComfyUI/custom_nodes', + ]; + return requiredDirs.includes(path); + }); + + const result = ComfyConfigManager.setUpComfyUI('/existing/ComfyUI'); + + expect(result).toBe('/existing/ComfyUI'); + }); + + it('should create ComfyUI subdirectory when it is missing', () => { + (fs.existsSync as jest.Mock).mockImplementationOnce((path: string) => { + if (path === '/some/base/path/ComfyUI') { + return false; + } + return true; + }); + + const result = ComfyConfigManager.setUpComfyUI('/some/base/path'); + + expect(result).toBe('/some/base/path/ComfyUI'); + }); + }); + + describe('isComfyUIDirectory', () => { + it('should return true when all required directories exist', () => { + (fs.existsSync as jest.Mock).mockImplementation((path: string) => { + const requiredDirs = [ + '/fake/path/models', + '/fake/path/input', + '/fake/path/user', + '/fake/path/output', + '/fake/path/custom_nodes', + ]; + return requiredDirs.includes(path); + }); + + const result = ComfyConfigManager.isComfyUIDirectory('/fake/path'); + + expect(result).toBe(true); + expect(fs.existsSync).toHaveBeenCalledTimes(5); + }); + + it('should return false when some required directories are missing', () => { + (fs.existsSync as jest.Mock) + .mockReturnValueOnce(true) // models exists + .mockReturnValueOnce(true) // input exists + .mockReturnValueOnce(false) // user missing + .mockReturnValueOnce(true) // output exists + .mockReturnValueOnce(true); // custom_nodes exists + + const result = ComfyConfigManager.isComfyUIDirectory('/fake/path'); + + expect(result).toBe(false); + }); + }); + + describe('createComfyDirectories', () => { + it('should create all necessary directories when none exist', () => { + (fs.existsSync as jest.Mock).mockReturnValue(false); + + ComfyConfigManager.createComfyDirectories('/fake/path/ComfyUI'); + + // Verify each required directory was created + expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/models', { recursive: true }); + expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/input', { recursive: true }); + expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/user', { recursive: true }); + expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/output', { recursive: true }); + expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/custom_nodes', { recursive: true }); + }); + }); + + describe('createComfyConfigFile', () => { + it('should create new config file when none exists', () => { + (fs.existsSync as jest.Mock).mockReturnValue(false); + + ComfyConfigManager.createComfyConfigFile('/fake/path', false); + + expect(fs.writeFileSync).toHaveBeenCalledTimes(1); + expect(fs.renameSync).not.toHaveBeenCalled(); + }); + + it('should backup existing config file when overwrite is true', () => { + (fs.existsSync as jest.Mock).mockImplementation((path: string) => { + return path === '/user/default/comfy.settings.json'; + }); + + ComfyConfigManager.createComfyConfigFile('/user/default', true); + + expect(fs.renameSync).toHaveBeenCalledTimes(1); + expect(fs.writeFileSync).toHaveBeenCalledTimes(1); + }); + + it('should handle backup failure gracefully', () => { + (fs.existsSync as jest.Mock).mockReturnValue(true); + (fs.renameSync as jest.Mock).mockImplementation(() => { + throw new Error('Backup failed'); + }); + + ComfyConfigManager.createComfyConfigFile('/fake/path', true); + + expect(fs.writeFileSync).not.toHaveBeenCalled(); + }); + }); + + describe('createNestedDirectories', () => { + it('should create nested directory structure correctly', () => { + (fs.existsSync as jest.Mock).mockReturnValue(false); + + const structure = ['dir1', ['dir2', ['subdir1', 'subdir2']], ['dir3', [['subdir3', ['subsubdir1']]]]]; + + ComfyConfigManager['createNestedDirectories']('/fake/path', structure); + + // Verify the correct paths were created + expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir1'), expect.any(Object)); + expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir2'), expect.any(Object)); + expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('subdir1'), expect.any(Object)); + expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('subsubdir1'), expect.any(Object)); + }); + + it('should handle invalid directory structure items', () => { + const invalidStructure = [ + 'dir1', + ['dir2'], // Invalid: array with only one item + [123, ['subdir1']], // Invalid: non-string directory name + ]; + + ComfyConfigManager['createNestedDirectories']('/fake/path', invalidStructure as DirectoryStructure); + + // Verify only valid directories were created + expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir1'), expect.any(Object)); + expect(fs.mkdirSync).not.toHaveBeenCalledWith(expect.stringContaining('subdir1'), expect.any(Object)); + }); + }); +}); diff --git a/src/config/comfyConfigManager.ts b/src/config/comfyConfigManager.ts new file mode 100644 index 00000000..0bbb0ba2 --- /dev/null +++ b/src/config/comfyConfigManager.ts @@ -0,0 +1,142 @@ +import fs from 'fs'; +import path from 'path'; +import log from 'electron-log/main'; + +export type DirectoryStructure = (string | DirectoryStructure)[]; + +export class ComfyConfigManager { + private static readonly DEFAULT_DIRECTORIES: DirectoryStructure = [ + 'custom_nodes', + 'input', + 'output', + ['user', ['default']], + [ + 'models', + [ + 'checkpoints', + 'clip', + 'clip_vision', + 'configs', + 'controlnet', + 'diffusers', + 'diffusion_models', + 'embeddings', + 'gligen', + 'hypernetworks', + 'loras', + 'photomaker', + 'style_models', + 'unet', + 'upscale_models', + 'vae', + 'vae_approx', + + // TODO(robinhuang): Remove when we have a better way to specify base model paths. + 'animatediff_models', + 'animatediff_motion_lora', + 'animatediff_video_formats', + 'liveportrait', + ['insightface', ['buffalo_1']], + ['blip', ['checkpoints']], + 'CogVideo', + ['xlabs', ['loras', 'controlnets']], + 'layerstyle', + 'LLM', + 'Joy_caption', + ], + ], + ]; + + private static readonly DEFAULT_CONFIG = { + 'Comfy.ColorPalette': 'dark', + 'Comfy.UseNewMenu': 'Top', + 'Comfy.Workflow.WorkflowTabsPosition': 'Topbar', + 'Comfy.Workflow.ShowMissingModelsWarning': true, + }; + + public static setUpComfyUI(localComfyDirectory: string): string { + if (!this.isComfyUIDirectory(localComfyDirectory)) { + log.info( + `Selected directory ${localComfyDirectory} is not a ComfyUI directory. Appending ComfyUI to install path.` + ); + localComfyDirectory = path.join(localComfyDirectory, 'ComfyUI'); + } + + this.createComfyDirectories(localComfyDirectory); + const userSettingsPath = path.join(localComfyDirectory, 'user', 'default'); + this.createComfyConfigFile(userSettingsPath, true); + return localComfyDirectory; + } + + public static createComfyConfigFile(userSettingsPath: string, overwrite: boolean = false): void { + const configFilePath = path.join(userSettingsPath, 'comfy.settings.json'); + + if (fs.existsSync(configFilePath) && overwrite) { + const backupFilePath = path.join(userSettingsPath, 'old_comfy.settings.json'); + try { + fs.renameSync(configFilePath, backupFilePath); + log.info(`Renaming existing user settings file to: ${backupFilePath}`); + } catch (error) { + log.error(`Failed to backup existing user settings file: ${error}`); + return; + } + } + + try { + fs.writeFileSync(configFilePath, JSON.stringify(this.DEFAULT_CONFIG, null, 2)); + log.info(`Created new ComfyUI config file at: ${configFilePath}`); + } catch (error) { + log.error(`Failed to create new ComfyUI config file: ${error}`); + } + } + + public static isComfyUIDirectory(directory: string): boolean { + const requiredSubdirs = ['models', 'input', 'user', 'output', 'custom_nodes']; + return requiredSubdirs.every((subdir) => fs.existsSync(path.join(directory, subdir))); + } + + static createComfyDirectories(localComfyDirectory: string): void { + log.info(`Creating ComfyUI directories in ${localComfyDirectory}`); + + try { + this.createNestedDirectories(localComfyDirectory, this.DEFAULT_DIRECTORIES); + } catch (error) { + log.error(`Failed to create ComfyUI directories: ${error}`); + } + } + + static createNestedDirectories(basePath: string, structure: DirectoryStructure): void { + structure.forEach((item) => { + if (typeof item === 'string') { + const dirPath = path.join(basePath, item); + this.createDirIfNotExists(dirPath); + } else if (Array.isArray(item) && item.length === 2) { + const [dirName, subDirs] = item; + if (typeof dirName === 'string') { + const newBasePath = path.join(basePath, dirName); + this.createDirIfNotExists(newBasePath); + if (Array.isArray(subDirs)) { + this.createNestedDirectories(newBasePath, subDirs); + } + } else { + log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`); + } + } else { + log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`); + } + }); + } + + /** + * Create a directory if not exists + * @param dirPath + */ + static createDirIfNotExists(dirPath: string): void { + if (!fs.existsSync(dirPath)) { + fs.mkdirSync(dirPath, { recursive: true }); + log.info(`Created directory: ${dirPath}`); + } else { + log.info(`Directory already exists: ${dirPath}`); + } + } +} diff --git a/src/main.ts b/src/main.ts index d260f6aa..69f5c772 100644 --- a/src/main.ts +++ b/src/main.ts @@ -10,7 +10,7 @@ import { IPCChannel, SENTRY_URL_ENDPOINT, } from './constants'; -import { app, BrowserWindow, dialog, screen, ipcMain, Menu, MenuItem, shell } from 'electron'; +import { app, BrowserWindow, dialog, screen, ipcMain, shell } from 'electron'; import log from 'electron-log/main'; import * as Sentry from '@sentry/electron/main'; import Store from 'electron-store'; @@ -24,6 +24,8 @@ 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'; dotenv.config(); @@ -344,37 +346,6 @@ function restartApp({ customMessage, delay }: { customMessage?: string; delay?: }); } -function buildMenu(): void { - const menu = Menu.getApplicationMenu(); - if (menu) { - const aboutMenuItem = { - label: 'About ComfyUI', - click: () => { - dialog.showMessageBox({ - title: 'About', - message: `ComfyUI v${app.getVersion()}`, - detail: 'Created by Comfy Org\nCopyright © 2024', - buttons: ['OK'], - }); - }, - }; - const helpMenuItem = menu.items.find((item) => item.role === 'help'); - if (helpMenuItem && helpMenuItem.submenu) { - helpMenuItem.submenu.append(new MenuItem(aboutMenuItem)); - Menu.setApplicationMenu(menu); - } else { - // If there's no Help menu, add one - menu.append( - new MenuItem({ - label: 'Help', - submenu: [aboutMenuItem], - }) - ); - Menu.setApplicationMenu(menu); - } - } -} - /** * Creates the main window. If the window already exists, it will return the existing window. * @param userResourcesPath The path to the user's resources. @@ -409,6 +380,7 @@ export const createWindow = async (userResourcesPath?: string): Promise { if (mainWindow) { @@ -709,131 +681,6 @@ const spawnPythonAsync = ( }); }; -function isComfyUIDirectory(directory: string): boolean { - const requiredSubdirs = ['models', 'input', 'user', 'output', 'custom_nodes']; - return requiredSubdirs.every((subdir) => fs.existsSync(path.join(directory, subdir))); -} - -type DirectoryStructure = (string | DirectoryStructure)[]; - -function createComfyDirectories(localComfyDirectory: string): void { - log.info(`Creating ComfyUI directories in ${localComfyDirectory}`); - - const directories: DirectoryStructure = [ - 'custom_nodes', - 'input', - 'output', - ['user', ['default']], - [ - 'models', - [ - 'checkpoints', - 'clip', - 'clip_vision', - 'configs', - 'controlnet', - 'diffusers', - 'diffusion_models', - 'embeddings', - 'gligen', - 'hypernetworks', - 'loras', - 'photomaker', - 'style_models', - 'unet', - 'upscale_models', - 'vae', - 'vae_approx', - - // TODO(robinhuang): Remove when we have a better way to specify base model paths. - 'animatediff_models', - 'animatediff_motion_lora', - 'animatediff_video_formats', - 'liveportrait', - ['insightface', ['buffalo_1']], - ['blip', ['checkpoints']], - 'CogVideo', - ['xlabs', ['loras', 'controlnets']], - 'layerstyle', - 'LLM', - 'Joy_caption', - ], - ], - ]; - try { - createNestedDirectories(localComfyDirectory, directories); - } catch (error) { - log.error(`Failed to create ComfyUI directories: ${error}`); - } - - const userSettingsPath = path.join(localComfyDirectory, 'user', 'default'); - createComfyConfigFile(userSettingsPath, true); -} - -function createNestedDirectories(basePath: string, structure: DirectoryStructure): void { - structure.forEach((item) => { - if (typeof item === 'string') { - const dirPath = path.join(basePath, item); - createDirIfNotExists(dirPath); - } else if (Array.isArray(item) && item.length === 2) { - const [dirName, subDirs] = item; - if (typeof dirName === 'string') { - const newBasePath = path.join(basePath, dirName); - createDirIfNotExists(newBasePath); - if (Array.isArray(subDirs)) { - createNestedDirectories(newBasePath, subDirs); - } - } else { - log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`); - } - } else { - log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`); - } - }); -} - -/** - * Create a directory if not exists - * @param dirPath - */ -function createDirIfNotExists(dirPath: string): void { - if (!fs.existsSync(dirPath)) { - fs.mkdirSync(dirPath, { recursive: true }); - log.info(`Created directory: ${dirPath}`); - } else { - log.info(`Directory already exists: ${dirPath}`); - } -} - -function createComfyConfigFile(userSettingsPath: string, overwrite: boolean = false): void { - const configContent: any = { - 'Comfy.ColorPalette': 'dark', - 'Comfy.UseNewMenu': 'Top', - 'Comfy.Workflow.WorkflowTabsPosition': 'Topbar', - 'Comfy.Workflow.ShowMissingModelsWarning': true, - }; - - const configFilePath = path.join(userSettingsPath, 'comfy.settings.json'); - - if (fs.existsSync(configFilePath) && overwrite) { - const backupFilePath = path.join(userSettingsPath, 'old_comfy.settings.json'); - try { - fs.renameSync(configFilePath, backupFilePath); - log.info(`Renaming existing user settings file to: ${backupFilePath}`); - } catch (error) { - log.error(`Failed to backup existing user settings file: ${error}`); - return; - } - } - - try { - fs.writeFileSync(configFilePath, JSON.stringify(configContent, null, 2)); - log.info(`Created new ComfyUI config file at: ${configFilePath}`); - } catch (error) { - log.error(`Failed to create new ComfyUI config file: ${error}`); - } -} - function findAvailablePort(startPort: number, endPort: number): Promise { return new Promise((resolve, reject) => { function tryPort(port: number) { @@ -881,18 +728,11 @@ async function handleFirstTimeSetup() { log.info('First time setup:', firstTimeSetup); if (firstTimeSetup) { sendRendererMessage(IPC_CHANNELS.SHOW_SELECT_DIRECTORY, null); - let selectedDirectory = await selectedInstallDirectory(); - if (!isComfyUIDirectory(selectedDirectory)) { - log.info( - `Selected directory ${selectedDirectory} is not a ComfyUI directory. Appending ComfyUI to install path.` - ); - selectedDirectory = path.join(selectedDirectory, 'ComfyUI'); - } - - createComfyDirectories(selectedDirectory); + const selectedDirectory = await selectedInstallDirectory(); + const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(selectedDirectory); const { modelConfigPath } = await determineResourcesPaths(); - await createModelConfigFiles(modelConfigPath, selectedDirectory); + await createModelConfigFiles(modelConfigPath, actualComfyDirectory); } else { sendRendererMessage(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null); } diff --git a/src/menu/menu.ts b/src/menu/menu.ts new file mode 100644 index 00000000..edfcde28 --- /dev/null +++ b/src/menu/menu.ts @@ -0,0 +1,32 @@ +import { app, dialog, Menu, MenuItem } from 'electron'; + +export function buildMenu(): void { + const menu = Menu.getApplicationMenu(); + if (menu) { + const aboutMenuItem = { + label: 'About ComfyUI', + click: () => { + dialog.showMessageBox({ + title: 'About', + message: `ComfyUI v${app.getVersion()}`, + detail: 'Created by Comfy Org\nCopyright © 2024', + buttons: ['OK'], + }); + }, + }; + const helpMenuItem = menu.items.find((item) => item.role === 'help'); + if (helpMenuItem && helpMenuItem.submenu) { + helpMenuItem.submenu.append(new MenuItem(aboutMenuItem)); + Menu.setApplicationMenu(menu); + } else { + // If there's no Help menu, add one + menu.append( + new MenuItem({ + label: 'Help', + submenu: [aboutMenuItem], + }) + ); + Menu.setApplicationMenu(menu); + } + } +}