From cc5baede02f6bff36ca2cbbd1b968ee2923d9f87 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Sun, 16 Mar 2025 17:40:04 +0000 Subject: [PATCH 1/5] Add support for importing model files --- package.json | 2 + src/constants.ts | 1 + src/handlers/importModelHandlers.ts | 63 +++++++++++++++++++++++++++++ src/main-process/comfyDesktopApp.ts | 2 + src/preload.ts | 19 ++++++++- yarn.lock | 11 ++++- 6 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 src/handlers/importModelHandlers.ts diff --git a/package.json b/package.json index f671e2ade..287de40d4 100644 --- a/package.json +++ b/package.json @@ -69,6 +69,7 @@ "@types/diff": "^7", "@types/electron-squirrel-startup": "^1.0.2", "@types/eslint__js": "^8.42.3", + "@types/mv": "^2", "@types/node": "^22.10.2", "@types/tar": "6.1.13", "@types/tmp": "^0.2.6", @@ -111,6 +112,7 @@ "electron-store": "8.2.0", "lodash": "4.17.21", "mixpanel": "^0.18.0", + "mv": "^2.1.1", "node-pty": "^1.0.0", "systeminformation": "^5.24.8", "tar": "^7.4.3", diff --git a/src/constants.ts b/src/constants.ts index 529555dd6..43cdb3d69 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -49,6 +49,7 @@ export const IPC_CHANNELS = { CAN_ACCESS_URL: 'can-access-url', START_TROUBLESHOOTING: 'start-troubleshooting', DISABLE_CUSTOM_NODES: 'disable-custom-nodes', + IMPORT_MODEL: 'import-model', } as const; export enum ProgressStatus { diff --git a/src/handlers/importModelHandlers.ts b/src/handlers/importModelHandlers.ts new file mode 100644 index 000000000..306266d4b --- /dev/null +++ b/src/handlers/importModelHandlers.ts @@ -0,0 +1,63 @@ +import { ipcMain } from 'electron'; +import log from 'electron-log/main'; +import { existsSync } from 'node:fs'; +import { copyFile, mkdir } from 'node:fs/promises'; +import mv from 'mv'; +import path from 'node:path'; + +import { ComfyServerConfig } from '@/config/comfyServerConfig'; + +import { IPC_CHANNELS } from '../constants'; + +export function registerImportModelHandlers() { + ipcMain.handle( + IPC_CHANNELS.IMPORT_MODEL, + async (_, filePath: string, type: string, mode: 'move' | 'copy'): Promise => { + try { + const configPath = ComfyServerConfig.configPath; + const config = await ComfyServerConfig.readConfigFile(configPath); + if (!config) { + throw new Error('Unable to read extra_model_paths.yaml'); + } + + // Find all config sections that have the type defined. + let sections = Object.entries(config) + .filter((c) => c[1][type]) + .sort((a, b) => +a[1].is_default - +b[1].is_default) + .map((c) => c[0]); + + // If no config sections with the type defined, use either the default or the first config section. + if (!sections.length) { + const defaultSection = Object.entries(config).find((c) => c[1].is_default)?.[0] ?? Object.keys(config)[0]; + sections = [defaultSection]; + } + + const targetConfig = config[sections[0]]; + let basePath = targetConfig.base_path; + if (targetConfig.download_model_base) { + basePath = path.join(basePath, targetConfig.download_model_base); + } + + const folderName = targetConfig[type] ?? type; + const destinationDir = path.join(basePath, folderName); + const destinationPath = path.join(destinationDir, path.basename(filePath)); + + if (!existsSync(destinationDir)) { + await mkdir(destinationDir, { recursive: true }); + } + + log.info(mode, filePath, '->', destinationPath); + + await (mode === 'move' ? new Promise((resolve, reject) => { + mv(filePath, destinationPath, (err) => { + if (err) reject(err instanceof Error ? err : new Error(String(err))); + else resolve(true); + }); + }) : copyFile(filePath, destinationPath)); + } catch (error) { + console.error(error); + throw error; + } + } + ); +} diff --git a/src/main-process/comfyDesktopApp.ts b/src/main-process/comfyDesktopApp.ts index f0c8b1eb3..0e6de8671 100644 --- a/src/main-process/comfyDesktopApp.ts +++ b/src/main-process/comfyDesktopApp.ts @@ -3,6 +3,7 @@ import { app, ipcMain } from 'electron'; import log from 'electron-log/main'; import { useComfySettings } from '@/config/comfySettings'; +import { registerImportModelHandlers } from '@/handlers/importModelHandlers'; import { DEFAULT_SERVER_ARGS, IPC_CHANNELS, ProgressStatus, ServerArgs } from '../constants'; import { DownloadManager } from '../models/DownloadManager'; @@ -76,6 +77,7 @@ export class ComfyDesktopApp implements HasTelemetry { registerIPCHandlers(): void { // Restart core ipcMain.handle(IPC_CHANNELS.RESTART_CORE, async (): Promise => await this.restartComfyServer()); + registerImportModelHandlers(); app.on('before-quit', () => { if (!this.comfyServer) return; diff --git a/src/preload.ts b/src/preload.ts index 5da78987a..319cc81a6 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -1,4 +1,4 @@ -import { contextBridge, ipcRenderer } from 'electron'; +import { contextBridge, ipcRenderer, webUtils } from 'electron'; import path from 'node:path'; import { DownloadStatus, ELECTRON_BRIDGE_API, IPC_CHANNELS, ProgressStatus } from './constants'; @@ -295,6 +295,23 @@ const electronAPI = { showContextMenu: (options?: ElectronContextMenuOptions): void => { return ipcRenderer.send(IPC_CHANNELS.SHOW_CONTEXT_MENU, options); }, + /** + * Get the path for a file. + * @param file The file to get the path for. + * @returns The path for the file. + */ + getFilePath: (file: File) => { + return webUtils.getPathForFile(file); + }, + /** + * Imports a model file. + * @param file The file to import. + * @param type The type of model to import. + * @param mode If the file should be moved or copied. + */ + importModel: (file: File, type: string, mode: 'move' | 'copy') => { + return ipcRenderer.invoke(IPC_CHANNELS.IMPORT_MODEL, webUtils.getPathForFile(file), type, mode); + }, Config: { /** * Finds the name of the last detected GPU type. Detection only runs during installation. diff --git a/yarn.lock b/yarn.lock index d041e6f07..b68cd828d 100644 --- a/yarn.lock +++ b/yarn.lock @@ -242,6 +242,7 @@ __metadata: "@types/electron-squirrel-startup": "npm:^1.0.2" "@types/eslint__js": "npm:^8.42.3" "@types/lodash": "npm:4.17.15" + "@types/mv": "npm:^2" "@types/node": "npm:^22.10.2" "@types/tar": "npm:6.1.13" "@types/tmp": "npm:^0.2.6" @@ -264,6 +265,7 @@ __metadata: lint-staged: "npm:^15.2.10" lodash: "npm:4.17.21" mixpanel: "npm:^0.18.0" + mv: "npm:^2.1.1" node-pty: "npm:^1.0.0" prettier: "npm:^3.3.3" rimraf: "npm:^6.0.1" @@ -2882,6 +2884,13 @@ __metadata: languageName: node linkType: hard +"@types/mv@npm:^2": + version: 2.1.4 + resolution: "@types/mv@npm:2.1.4" + checksum: 10c0/9372401a4381aacc0dd004b8ba37c954a26c98b1937952cb0dc9c8933c904dceb509e2e697c1f9d21c62eb0d4de965f669d157ad1e4a8c8ea1d28252c26f8faa + languageName: node + linkType: hard + "@types/mysql@npm:2.15.26": version: 2.15.26 resolution: "@types/mysql@npm:2.15.26" @@ -8895,7 +8904,7 @@ __metadata: languageName: node linkType: hard -"mv@npm:~2": +"mv@npm:^2.1.1, mv@npm:~2": version: 2.1.1 resolution: "mv@npm:2.1.1" dependencies: From 27d6b7d7f37a478171d9226f3f06c1fad6bd3a5a Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Sun, 16 Mar 2025 19:12:58 +0000 Subject: [PATCH 2/5] Format --- src/handlers/importModelHandlers.ts | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/handlers/importModelHandlers.ts b/src/handlers/importModelHandlers.ts index 306266d4b..9fc31205d 100644 --- a/src/handlers/importModelHandlers.ts +++ b/src/handlers/importModelHandlers.ts @@ -1,8 +1,8 @@ import { ipcMain } from 'electron'; import log from 'electron-log/main'; +import mv from 'mv'; import { existsSync } from 'node:fs'; import { copyFile, mkdir } from 'node:fs/promises'; -import mv from 'mv'; import path from 'node:path'; import { ComfyServerConfig } from '@/config/comfyServerConfig'; @@ -48,12 +48,14 @@ export function registerImportModelHandlers() { log.info(mode, filePath, '->', destinationPath); - await (mode === 'move' ? new Promise((resolve, reject) => { - mv(filePath, destinationPath, (err) => { - if (err) reject(err instanceof Error ? err : new Error(String(err))); - else resolve(true); - }); - }) : copyFile(filePath, destinationPath)); + await (mode === 'move' + ? new Promise((resolve, reject) => { + mv(filePath, destinationPath, (err) => { + if (err) reject(err instanceof Error ? err : new Error(String(err))); + else resolve(true); + }); + }) + : copyFile(filePath, destinationPath)); } catch (error) { console.error(error); throw error; From 06f6b7a4c560e164aab616f3959b062a632d7ab9 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Thu, 27 Mar 2025 10:54:49 +0800 Subject: [PATCH 3/5] Simplify import by using config basepath --- src/handlers/importModelHandlers.ts | 29 +++++------------------------ 1 file changed, 5 insertions(+), 24 deletions(-) diff --git a/src/handlers/importModelHandlers.ts b/src/handlers/importModelHandlers.ts index 9fc31205d..641920f93 100644 --- a/src/handlers/importModelHandlers.ts +++ b/src/handlers/importModelHandlers.ts @@ -5,7 +5,7 @@ import { existsSync } from 'node:fs'; import { copyFile, mkdir } from 'node:fs/promises'; import path from 'node:path'; -import { ComfyServerConfig } from '@/config/comfyServerConfig'; +import { useDesktopConfig } from '@/store/desktopConfig'; import { IPC_CHANNELS } from '../constants'; @@ -14,32 +14,13 @@ export function registerImportModelHandlers() { IPC_CHANNELS.IMPORT_MODEL, async (_, filePath: string, type: string, mode: 'move' | 'copy'): Promise => { try { - const configPath = ComfyServerConfig.configPath; - const config = await ComfyServerConfig.readConfigFile(configPath); - if (!config) { - throw new Error('Unable to read extra_model_paths.yaml'); - } - - // Find all config sections that have the type defined. - let sections = Object.entries(config) - .filter((c) => c[1][type]) - .sort((a, b) => +a[1].is_default - +b[1].is_default) - .map((c) => c[0]); - - // If no config sections with the type defined, use either the default or the first config section. - if (!sections.length) { - const defaultSection = Object.entries(config).find((c) => c[1].is_default)?.[0] ?? Object.keys(config)[0]; - sections = [defaultSection]; - } + const basePath = useDesktopConfig().get('basePath'); - const targetConfig = config[sections[0]]; - let basePath = targetConfig.base_path; - if (targetConfig.download_model_base) { - basePath = path.join(basePath, targetConfig.download_model_base); + if (!basePath) { + throw new Error('Base path is not set'); } - const folderName = targetConfig[type] ?? type; - const destinationDir = path.join(basePath, folderName); + const destinationDir = path.join(basePath, 'models', type); const destinationPath = path.join(destinationDir, path.basename(filePath)); if (!existsSync(destinationDir)) { From 8a3e64c0bd8522fcad65bc2c34aaea8d8032afdc Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Thu, 27 Mar 2025 10:57:39 +0800 Subject: [PATCH 4/5] formatting --- src/handlers/importModelHandlers.ts | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/handlers/importModelHandlers.ts b/src/handlers/importModelHandlers.ts index 641920f93..708f08274 100644 --- a/src/handlers/importModelHandlers.ts +++ b/src/handlers/importModelHandlers.ts @@ -29,14 +29,12 @@ export function registerImportModelHandlers() { log.info(mode, filePath, '->', destinationPath); - await (mode === 'move' - ? new Promise((resolve, reject) => { - mv(filePath, destinationPath, (err) => { - if (err) reject(err instanceof Error ? err : new Error(String(err))); - else resolve(true); - }); - }) - : copyFile(filePath, destinationPath)); + await (mode == 'move' ? new Promise((resolve, reject) => { + mv(filePath, destinationPath, (err) => { + if (err) reject(err instanceof Error ? err : new Error(String(err))); + else resolve(true); + }); + }) : copyFile(filePath, destinationPath)); } catch (error) { console.error(error); throw error; From c8635b69470d4bfc8c2985b84a0df0ca831e9fd0 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Thu, 27 Mar 2025 11:08:10 +0800 Subject: [PATCH 5/5] format --- src/handlers/importModelHandlers.ts | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/handlers/importModelHandlers.ts b/src/handlers/importModelHandlers.ts index 708f08274..5030ab2ca 100644 --- a/src/handlers/importModelHandlers.ts +++ b/src/handlers/importModelHandlers.ts @@ -29,12 +29,14 @@ export function registerImportModelHandlers() { log.info(mode, filePath, '->', destinationPath); - await (mode == 'move' ? new Promise((resolve, reject) => { - mv(filePath, destinationPath, (err) => { - if (err) reject(err instanceof Error ? err : new Error(String(err))); - else resolve(true); - }); - }) : copyFile(filePath, destinationPath)); + await (mode == 'move' + ? new Promise((resolve, reject) => { + mv(filePath, destinationPath, (err) => { + if (err) reject(err instanceof Error ? err : new Error(String(err))); + else resolve(true); + }); + }) + : copyFile(filePath, destinationPath)); } catch (error) { console.error(error); throw error;