diff --git a/src/main-process/comfyInstallation.ts b/src/main-process/comfyInstallation.ts index d2e5e9627..0da7e7745 100644 --- a/src/main-process/comfyInstallation.ts +++ b/src/main-process/comfyInstallation.ts @@ -44,7 +44,7 @@ export class ComfyInstallation { /** `true` if Manager needs toml and uv to be installed, otherwise `false`. */ get needsRequirementsUpdate() { - return this.validation.upgradePackages === 'warning'; + return this.validation.needsPackageUpdate === true; } /** @@ -184,12 +184,14 @@ export class ComfyInstallation { // Python packages try { const result = await venv.hasRequirements(); - if (result === 'package-upgrade') { - validation.pythonPackages = 'OK'; - validation.upgradePackages = 'warning'; - } else { - validation.pythonPackages = result; - if (result !== 'OK') log.error('Virtual environment is incomplete.'); + switch (result.status) { + case 'upgrade': + validation.pythonPackages = 'OK'; + validation.needsPackageUpdate = true; + break; + case 'ok': + validation.pythonPackages = 'OK'; + break; } } catch (error) { log.error('Failed to read venv packages.', error); diff --git a/src/preload.ts b/src/preload.ts index f795f38e8..478cd0f55 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -118,7 +118,8 @@ export interface InstallValidation { uv?: ValidationIssueState; git?: ValidationIssueState; vcRedist?: ValidationIssueState; - upgradePackages?: ValidationIssueState; + /** True if packages need updating (not an error, will auto-update) */ + needsPackageUpdate?: boolean; } const electronAPI = { diff --git a/src/virtualEnvironment.ts b/src/virtualEnvironment.ts index c626cc04c..b2701d5e7 100644 --- a/src/virtualEnvironment.ts +++ b/src/virtualEnvironment.ts @@ -58,6 +58,15 @@ type TorchPackageVersions = Record; const TORCH_PACKAGE_NAMES: TorchPackageName[] = ['torch', 'torchaudio', 'torchvision']; +export type RequirementsCheckStatus = 'ok' | 'upgrade'; + +export type RequirementsCheckReason = 'requirements-diff' | 'nvidia-torch'; + +export type RequirementsCheckResult = { + status: RequirementsCheckStatus; + reason?: RequirementsCheckReason; +}; + export function getPipInstallArgs(config: PipInstallConfig): string[] { const installArgs = ['pip', 'install']; @@ -327,7 +336,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor { const requirementsStatus = await this.hasRequirements(); - if (requirementsStatus === 'OK') { + if (requirementsStatus.status === 'ok') { log.info('Skipping requirements installation - all requirements already installed'); } else { log.info('Starting manual install - venv missing requirements'); @@ -888,11 +897,9 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor { * Checks if the virtual environment has all the required packages of ComfyUI core. * * Parses the text output of `uv pip install --dry-run -r requirements.txt`. - * @returns `'OK'` if pip install does not detect any missing packages, - * `'manager-upgrade'` if `uv` and `toml` are missing, - * or `'error'` when any other combination of packages are missing. + * @returns Result describing whether requirements are satisfied or require an upgrade. */ - async hasRequirements(): Promise<'OK' | 'error' | 'package-upgrade'> { + async hasRequirements(): Promise { const checkRequirements = async (requirementsPath: string) => { const args = ['pip', 'install', '--dry-run', '-r', requirementsPath]; log.info(`Running uv command directly: ${args.join(' ')}`); @@ -918,40 +925,6 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor { return venvOk; }; - // Manager upgrade in 0.4.18 - uv, toml (exactly) - const isManagerUpgrade = (output: string) => { - // Match the original case: 2 packages (uv + toml) | Added in https://github.com/ltdrdata/ComfyUI-Manager/commit/816a53a7b1a057af373c458ebf80aaae565b996b - // Match the new case: 1 package (chardet) | Added in https://github.com/ltdrdata/ComfyUI-Manager/commit/60a5e4f2614c688b41a1ebaf0694953eb26db38a - const anyCombination = /\bWould install [1-3] packages?(\s+\+ (toml|uv|chardet)==[\d.]+){1,3}\s*$/; - return anyCombination.test(output); - }; - - // Package upgrade in 0.4.21 - aiohttp, av, yarl - const isCoreUpgrade = (output: string) => { - const lines = output.split('\n'); - let adds = 0; - for (const line of lines) { - // Reject upgrade if removing an unrecognised package - if ( - line.search( - /^\s*- (?!aiohttp|av|yarl|comfyui-workflow-templates|comfyui-embedded-docs|pydantic|pydantic-core|pydantic-settings|annotated-types|typing-inspection|alembic|sqlalchemy|greenlet|mako|python-dotenv).*==/ - ) !== -1 - ) - return false; - if (line.search(/^\s*\+ /) !== -1) { - if ( - line.search( - /^\s*\+ (aiohttp|av|yarl|comfyui-workflow-templates|comfyui-embedded-docs|pydantic|pydantic-core|pydantic-settings|annotated-types|typing-inspection|alembic|sqlalchemy|greenlet|mako|python-dotenv)==/ - ) === -1 - ) - return false; - adds++; - } - // An unexpected package means this is not a package upgrade - } - return adds > 0; - }; - const coreOutput = await checkRequirements(this.comfyUIRequirementsPath); if (!(await pathAccessible(this.comfyUIManagerRequirementsPath))) { throw new Error( @@ -964,31 +937,21 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor { const coreOk = hasAllPackages(coreOutput); const managerOk = hasAllPackages(managerOutput); - const upgradeCore = !coreOk && isCoreUpgrade(coreOutput); - const upgradeManager = !managerOk && isManagerUpgrade(managerOutput); - - if ((managerOk && upgradeCore) || (coreOk && upgradeManager) || (upgradeCore && upgradeManager)) { - log.info('Package update of known packages required. Core:', upgradeCore, 'Manager:', upgradeManager); - return 'package-upgrade'; - } - if (!coreOk || !managerOk) { - log.info('Requirements are out of date. Treating as package upgrade.', { + log.info('Requirements out of date. Scheduling package update.', { coreOk, managerOk, - upgradeCore, - upgradeManager, }); - return 'package-upgrade'; + return { status: 'upgrade', reason: 'requirements-diff' }; } if (await this.needsNvidiaTorchUpgrade()) { log.info('NVIDIA PyTorch version out of date. Treating as package upgrade.'); - return 'package-upgrade'; + return { status: 'upgrade', reason: 'nvidia-torch' }; } - log.debug('hasRequirements result:', 'OK'); - return 'OK'; + log.debug('hasRequirements result:', 'ok'); + return { status: 'ok' }; } /** diff --git a/tests/integration/post-install/troubleshootingVenv.spec.ts b/tests/integration/post-install/troubleshootingVenv.spec.ts index 2fdbdd028..ab0848402 100644 --- a/tests/integration/post-install/troubleshootingVenv.spec.ts +++ b/tests/integration/post-install/troubleshootingVenv.spec.ts @@ -15,18 +15,13 @@ test.describe('Troubleshooting - broken venv', () => { test.slow(); await troubleshooting.expectReady(); - const { resetVenvCard, installPythonPackagesCard } = troubleshooting; + const { resetVenvCard } = troubleshooting; await expect(resetVenvCard.rootEl).toBeVisible(); await resetVenvCard.button.click(); await troubleshooting.confirmRecreateVenvButton.click(); await expect(resetVenvCard.isRunningIndicator).toBeVisible(); - await expect(installPythonPackagesCard.rootEl).toBeVisible({ timeout: 60 * 1000 }); - await installPythonPackagesCard.button.click(); - await troubleshooting.confirmInstallPythonPackagesButton.click(); - await expect(installPythonPackagesCard.isRunningIndicator).toBeVisible(); - // Venv fixed - server should start await installedApp.waitUntilLoaded(3 * 60 * 1000); }); diff --git a/tests/unit/install/installationManager.test.ts b/tests/unit/install/installationManager.test.ts index af9f92969..a7d71a70a 100644 --- a/tests/unit/install/installationManager.test.ts +++ b/tests/unit/install/installationManager.test.ts @@ -86,7 +86,7 @@ vi.mock('@/virtualEnvironment', () => { return { VirtualEnvironment: vi.fn(() => ({ exists: vi.fn(() => Promise.resolve(true)), - hasRequirements: vi.fn(() => Promise.resolve(true)), + hasRequirements: vi.fn(() => Promise.resolve({ status: 'ok' })), pythonInterpreterPath: 'valid/python', uvPath: 'valid/uv', venvPath: 'valid/venv', diff --git a/tests/unit/main-process/comfyServer.test.ts b/tests/unit/main-process/comfyServer.test.ts index 6f056f788..c86623e87 100644 --- a/tests/unit/main-process/comfyServer.test.ts +++ b/tests/unit/main-process/comfyServer.test.ts @@ -64,13 +64,7 @@ const test = baseTest.extend({ transports: { file: { transforms: [] } }, } as unknown as MainLogger & { default: MainLogger }); - const server = new ComfyServer( - basePath, - mockServerArgs, - mockVirtualEnvironment as any, - mockAppWindow as any, - mockTelemetry as any - ); + const server = new ComfyServer(basePath, mockServerArgs, mockVirtualEnvironment, mockAppWindow, mockTelemetry); await use(server); }, runningServer: async ({ server, mockProcess }, use) => { diff --git a/tests/unit/virtualEnvironment.test.ts b/tests/unit/virtualEnvironment.test.ts index 2e6af3a4d..dbc20a426 100644 --- a/tests/unit/virtualEnvironment.test.ts +++ b/tests/unit/virtualEnvironment.test.ts @@ -146,8 +146,8 @@ test.for(allCombinations)('hasRequirements', async ({ core, manager }, { virtual mockSpawnForPackages(core); mockSpawnForPackages(manager); - const result = core.length + manager.length === 0 ? 'OK' : 'package-upgrade'; - await expect(virtualEnv.hasRequirements()).resolves.toBe(result); + const result = core.length + manager.length === 0 ? 'ok' : 'upgrade'; + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: result }); expect(log.info).toHaveBeenCalledWith(expect.stringContaining('pip install --dry-run -r')); }); @@ -184,60 +184,59 @@ describe('VirtualEnvironment', () => { mockSpawnOutputOnce('Would make no changes\n'); mockSpawnOutputOnce('Would make no changes\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('OK'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'ok' }); expect(log.info).toHaveBeenCalledWith(expect.stringContaining('pip install --dry-run -r')); }); - test('returns package-upgrade when packages are missing and not a known upgrade case', async ({ virtualEnv }) => { + test('returns upgrade when packages are missing', async ({ virtualEnv }) => { mockSpawnOutputOnce(' + unknown_package==1.0.0\n'); mockSpawnOutputOnce('Would make no changes\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); expect(log.info).toHaveBeenCalledWith( - expect.stringContaining('Requirements are out of date. Treating as package upgrade.'), - expect.objectContaining({ coreOk: false, managerOk: true, upgradeCore: false, upgradeManager: false }) + expect.stringContaining('Requirements out of date. Scheduling package update.'), + expect.objectContaining({ coreOk: false, managerOk: true }) ); }); - test('returns package-upgrade for manager upgrade case', async ({ virtualEnv }) => { + test('returns upgrade when manager requirements differ (single package)', async ({ virtualEnv }) => { mockSpawnOutputOnce('Would make no changes\n'); mockSpawnOutputOnce('Would install 1 package \n + chardet==5.2.0\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); expect(log.info).toHaveBeenCalledWith( - 'Package update of known packages required. Core:', - false, - 'Manager:', - true + expect.stringContaining('Requirements out of date. Scheduling package update.'), + expect.objectContaining({ coreOk: true, managerOk: false }) ); }); - test('returns package-upgrade for manager upgrade case', async ({ virtualEnv }) => { + test('returns upgrade when manager requirements differ (multiple packages)', async ({ virtualEnv }) => { mockSpawnOutputOnce('Would make no changes\n'); mockSpawnOutputOnce('Would install 2 packages \n + uv==1.0.0 \n + toml==1.0.0\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); expect(log.info).toHaveBeenCalledWith( - 'Package update of known packages required. Core:', - false, - 'Manager:', - true + expect.stringContaining('Requirements out of date. Scheduling package update.'), + expect.objectContaining({ coreOk: true, managerOk: false }) ); }); - test('returns package-upgrade for core + manager upgrade case', async ({ virtualEnv }) => { + test('returns upgrade when core and manager requirements differ', async ({ virtualEnv }) => { mockSpawnOutputOnce('Would install 3 packages \n + av==1.0.0 \n + yarl==12.0.8 \n + aiohttp==3.9.0\n'); mockSpawnOutputOnce('Would install 2 packages \n + uv==1.0.0 \n + toml==1.0.0\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); - expect(log.info).toHaveBeenCalledWith('Package update of known packages required. Core:', true, 'Manager:', true); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); + expect(log.info).toHaveBeenCalledWith( + expect.stringContaining('Requirements out of date. Scheduling package update.'), + expect.objectContaining({ coreOk: false, managerOk: false }) + ); }); - test('returns package-upgrade for core upgrade case', async ({ virtualEnv }) => { + test('returns upgrade when core requirements differ', async ({ virtualEnv }) => { mockSpawnOutputOnce('Would install 1 package \n + av==1.0.0\n'); mockSpawnOutputOnce('Would make no changes\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); }); test('throws error when pip command fails', async ({ virtualEnv }) => { @@ -256,14 +255,42 @@ describe('VirtualEnvironment', () => { mockSpawnOutputOnce('', 0, null, 'Would make no changes\n'); mockSpawnOutputOnce('', 0, null, 'Would make no changes\n'); - await expect(virtualEnv.hasRequirements()).resolves.toBe('OK'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'ok' }); }); - test('rejects core upgrade with unrecognized package removal', async ({ virtualEnv }) => { + test('returns upgrade when requirements remove packages', async ({ virtualEnv }) => { mockSpawnOutputOnce(' - unknown-package==1.0.0\n + aiohttp==3.9.0\n', 0, null); mockSpawnOutputOnce('Would make no changes\n', 0, null); - await expect(virtualEnv.hasRequirements()).resolves.toBe('package-upgrade'); + await expect(virtualEnv.hasRequirements()).resolves.toMatchObject({ status: 'upgrade' }); + }); + + test('returns upgrade with nvidia-torch reason when torch packages are outdated', async () => { + vi.stubGlobal('process', { + ...process, + resourcesPath: path.join(__dirname, '../resources'), + }); + + const nvidiaEnv = new VirtualEnvironment('/mock/venv', { + telemetry: mockTelemetry, + selectedDevice: 'nvidia', + pythonVersion: '3.12', + }); + + mockSpawnOutputOnce('Would make no changes\n'); + mockSpawnOutputOnce('Would make no changes\n'); + mockSpawnOutputOnce( + JSON.stringify([ + { name: 'torch', version: '2.0.0+cu121' }, + { name: 'torchaudio', version: '2.0.0+cu121' }, + { name: 'torchvision', version: '0.15.0+cu121' }, + ]) + ); + + await expect(nvidiaEnv.hasRequirements()).resolves.toMatchObject({ + status: 'upgrade', + reason: 'nvidia-torch', + }); }); });