From d32dc35b6dfd87abaf07d2527e341148c0455272 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian <157918347+Xiaohan-Tian@users.noreply.github.com> Date: Mon, 11 May 2026 22:27:22 -0700 Subject: [PATCH 01/19] feat: added local separator implementation --- package-lock.json | 139 ++++- package.json | 1 + src/components/KGOnePanel.css | 95 ++- src/components/KGOnePanel.test.tsx | 180 ++++++ src/components/KGOnePanel.tsx | 396 ++++++++++++- src/components/Toolbar.tsx | 6 +- src/core/io/LocalSeparatorModelCache.test.ts | 151 +++++ src/util/localSeparatorConfig.ts | 31 + src/util/localSeparatorCpuDsp.ts | 139 +++++ src/util/localSeparatorGpuDsp.ts | 220 +++++++ src/util/localSeparatorModelCache.ts | 126 ++++ src/util/localSeparatorRunner.ts | 589 +++++++++++++++++++ src/util/localSeparatorRuntime.ts | 95 +++ src/util/localSeparatorShared.ts | 275 +++++++++ src/util/localSeparatorTiming.ts | 52 ++ src/util/localSeparatorTypes.ts | 49 ++ 16 files changed, 2506 insertions(+), 38 deletions(-) create mode 100644 src/components/KGOnePanel.test.tsx create mode 100644 src/core/io/LocalSeparatorModelCache.test.ts create mode 100644 src/util/localSeparatorConfig.ts create mode 100644 src/util/localSeparatorCpuDsp.ts create mode 100644 src/util/localSeparatorGpuDsp.ts create mode 100644 src/util/localSeparatorModelCache.ts create mode 100644 src/util/localSeparatorRunner.ts create mode 100644 src/util/localSeparatorRuntime.ts create mode 100644 src/util/localSeparatorShared.ts create mode 100644 src/util/localSeparatorTiming.ts create mode 100644 src/util/localSeparatorTypes.ts diff --git a/package-lock.json b/package-lock.json index 3ce903b..bc951f5 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,18 +1,19 @@ { "name": "K.G.Studio", - "version": "0.15.0-build.20260510", + "version": "0.16.0-build.20260510", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "K.G.Studio", - "version": "0.15.0-build.20260510", + "version": "0.16.0-build.20260510", "dependencies": { "@breezystack/lamejs": "^1.2.7", "class-transformer": "^0.5.1", "fft.js": "^4.0.4", "idb": "^8.0.3", "jszip": "^3.10.1", + "onnxruntime-web": "^1.26.0", "openai": "^6.33.0", "react": "^19.1.0", "react-dom": "^19.1.0", @@ -1628,6 +1629,70 @@ "dev": true, "license": "MIT" }, + "node_modules/@protobufjs/aspromise": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz", + "integrity": "sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/base64": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/@protobufjs/base64/-/base64-1.1.2.tgz", + "integrity": "sha512-AZkcAA5vnN/v4PDqKyMR5lx7hZttPDgClv83E//FMNhR2TMcLUhfRUBHCmSl0oi9zMgDDqRUJkSxO3wm85+XLg==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/codegen": { + "version": "2.0.5", + "resolved": "https://registry.npmjs.org/@protobufjs/codegen/-/codegen-2.0.5.tgz", + "integrity": "sha512-zgXFLzW3Ap33e6d0Wlj4MGIm6Ce8O89n/apUaGNB/jx+hw+ruWEp7EwGUshdLKVRCxZW12fp9r40E1mQrf/34g==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/eventemitter": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@protobufjs/eventemitter/-/eventemitter-1.1.0.tgz", + "integrity": "sha512-j9ednRT81vYJ9OfVuXG6ERSTdEL1xVsNgqpkxMsbIabzSo3goCjDIveeGv5d03om39ML71RdmrGNjG5SReBP/Q==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/fetch": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@protobufjs/fetch/-/fetch-1.1.0.tgz", + "integrity": "sha512-lljVXpqXebpsijW71PZaCYeIcE5on1w5DlQy5WH6GLbFryLUrBD4932W/E2BSpfRJWseIL4v/KPgBFxDOIdKpQ==", + "license": "BSD-3-Clause", + "dependencies": { + "@protobufjs/aspromise": "^1.1.1", + "@protobufjs/inquire": "^1.1.0" + } + }, + "node_modules/@protobufjs/float": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/@protobufjs/float/-/float-1.0.2.tgz", + "integrity": "sha512-Ddb+kVXlXst9d+R9PfTIxh1EdNkgoRe5tOX6t01f1lYWOvJnSPDBlG241QLzcyPdoNTsblLUdujGSE4RzrTZGQ==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/inquire": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@protobufjs/inquire/-/inquire-1.1.1.tgz", + "integrity": "sha512-mnzgDV26ueAvk7rsbt9L7bE0SuAoqyuys/sMMrmVcN5x9VsxpcG3rqAUSgDyLp0UZlmNfIbQ4fHfCtreVBk8Ew==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/path": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/@protobufjs/path/-/path-1.1.2.tgz", + "integrity": "sha512-6JOcJ5Tm08dOHAbdR3GrvP+yUUfkjG5ePsHYczMFLq3ZmMkAD98cDgcT2iA1lJ9NVwFd4tH/iSSoe44YWkltEA==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/pool": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@protobufjs/pool/-/pool-1.1.0.tgz", + "integrity": "sha512-0kELaGSIDBKvcgS4zkjz1PeddatrjYcmMWOlAuAPwAeccUrPHdUqo/J6LiymHHEiJT5NrF1UVwxY14f+fy4WQw==", + "license": "BSD-3-Clause" + }, + "node_modules/@protobufjs/utf8": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@protobufjs/utf8/-/utf8-1.1.1.tgz", + "integrity": "sha512-oOAWABowe8EAbMyWKM0tYDKi8Yaox52D+HWZhAIJqQXbqe0xI/GV7FhLWqlEKreMkfDjshR5FKgi3mnle0h6Eg==", + "license": "BSD-3-Clause" + }, "node_modules/@rolldown/pluginutils": { "version": "1.0.0-beta.19", "resolved": "https://registry.npmjs.org/@rolldown/pluginutils/-/pluginutils-1.0.0-beta.19.tgz", @@ -2630,7 +2695,6 @@ "version": "22.18.0", "resolved": "https://registry.npmjs.org/@types/node/-/node-22.18.0.tgz", "integrity": "sha512-m5ObIqwsUp6BZzyiy4RdZpzWGub9bqLJMvZDD0QMXhxjqMHMENlj+SqF5QxoUwaQNFe+8kz8XM8ZQhqkQPTgMQ==", - "dev": true, "license": "MIT", "dependencies": { "undici-types": "~6.21.0" @@ -4802,6 +4866,12 @@ "node": ">=16" } }, + "node_modules/flatbuffers": { + "version": "25.9.23", + "resolved": "https://registry.npmjs.org/flatbuffers/-/flatbuffers-25.9.23.tgz", + "integrity": "sha512-MI1qs7Lo4Syw0EOzUl0xjs2lsoeqFku44KpngfIduHBYvzm8h2+7K8YMQh1JtVVVrUvhLpNwqVi4DERegUJhPQ==", + "license": "Apache-2.0" + }, "node_modules/flatted": { "version": "3.3.3", "resolved": "https://registry.npmjs.org/flatted/-/flatted-3.3.3.tgz", @@ -5087,6 +5157,12 @@ "dev": true, "license": "MIT" }, + "node_modules/guid-typescript": { + "version": "1.0.9", + "resolved": "https://registry.npmjs.org/guid-typescript/-/guid-typescript-1.0.9.tgz", + "integrity": "sha512-Y8T4vYhEfwJOTbouREvG+3XDsjr8E3kIr7uf+JZ0BYloFsttiHU0WfvANVsR7TxNUJa/WpCnw/Ino/p+DeBhBQ==", + "license": "ISC" + }, "node_modules/handlebars": { "version": "4.7.8", "resolved": "https://registry.npmjs.org/handlebars/-/handlebars-4.7.8.tgz", @@ -6121,6 +6197,12 @@ "dev": true, "license": "MIT" }, + "node_modules/long": { + "version": "5.3.2", + "resolved": "https://registry.npmjs.org/long/-/long-5.3.2.tgz", + "integrity": "sha512-mNAgZ1GmyNhD7AuqnTG3/VQ26o760+ZYBPKjPvugO8+nLbYfX6TVpJPseBvopbdY+qpZ/lKUnmEc1LeZYS3QAA==", + "license": "Apache-2.0" + }, "node_modules/longest-streak": { "version": "3.1.0", "resolved": "https://registry.npmjs.org/longest-streak/-/longest-streak-3.1.0.tgz", @@ -10920,6 +11002,26 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/onnxruntime-common": { + "version": "1.26.0", + "resolved": "https://registry.npmjs.org/onnxruntime-common/-/onnxruntime-common-1.26.0.tgz", + "integrity": "sha512-qVyMR4lcWgbkc4getFV+GQijsTnbg/siteoqcDwa3sI/LxbrMSNw4ePyvCq/ymdQaRomCA7YuWmhzsswxvymdw==", + "license": "MIT" + }, + "node_modules/onnxruntime-web": { + "version": "1.26.0", + "resolved": "https://registry.npmjs.org/onnxruntime-web/-/onnxruntime-web-1.26.0.tgz", + "integrity": "sha512-LbRr/8zZt2xilI2smrVQGGKINo0U46i8qJp+UXyMBGfqN7KjnH1BiwCwLwyNIVV4i9CKFv7Sf4PwLKWnT8/bEA==", + "license": "MIT", + "dependencies": { + "flatbuffers": "^25.1.24", + "guid-typescript": "^1.0.9", + "long": "^5.2.3", + "onnxruntime-common": "1.26.0", + "platform": "^1.3.6", + "protobufjs": "^7.2.4" + } + }, "node_modules/openai": { "version": "6.33.0", "resolved": "https://registry.npmjs.org/openai/-/openai-6.33.0.tgz", @@ -11341,6 +11443,12 @@ "node": ">=4" } }, + "node_modules/platform": { + "version": "1.3.6", + "resolved": "https://registry.npmjs.org/platform/-/platform-1.3.6.tgz", + "integrity": "sha512-fnWVljUchTro6RiCFvCXBbNhJc2NijN7oIQxbwsyL0buWJPG85v81ehlHI9fXrJsMNgTofEoWIQeClKpgxFLrg==", + "license": "MIT" + }, "node_modules/postcss": { "version": "8.5.6", "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.6.tgz", @@ -11442,6 +11550,30 @@ "dev": true, "license": "ISC" }, + "node_modules/protobufjs": { + "version": "7.5.7", + "resolved": "https://registry.npmjs.org/protobufjs/-/protobufjs-7.5.7.tgz", + "integrity": "sha512-NGnrxS/nLKUo5nkbVQxlC71sB4hdfImdYIbFeSCidxtwATx0AHRPcANSLd0q5Bb2BkoSWo2iisQhGg5/r+ihbA==", + "hasInstallScript": true, + "license": "BSD-3-Clause", + "dependencies": { + "@protobufjs/aspromise": "^1.1.2", + "@protobufjs/base64": "^1.1.2", + "@protobufjs/codegen": "^2.0.5", + "@protobufjs/eventemitter": "^1.1.0", + "@protobufjs/fetch": "^1.1.0", + "@protobufjs/float": "^1.0.2", + "@protobufjs/inquire": "^1.1.1", + "@protobufjs/path": "^1.1.2", + "@protobufjs/pool": "^1.1.0", + "@protobufjs/utf8": "^1.1.1", + "@types/node": ">=13.7.0", + "long": "^5.0.0" + }, + "engines": { + "node": ">=12.0.0" + } + }, "node_modules/punycode": { "version": "2.3.1", "resolved": "https://registry.npmjs.org/punycode/-/punycode-2.3.1.tgz", @@ -13531,7 +13663,6 @@ "version": "6.21.0", "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", - "dev": true, "license": "MIT" }, "node_modules/unicorn-magic": { diff --git a/package.json b/package.json index 8ab1ffb..68de2aa 100644 --- a/package.json +++ b/package.json @@ -19,6 +19,7 @@ "fft.js": "^4.0.4", "idb": "^8.0.3", "jszip": "^3.10.1", + "onnxruntime-web": "^1.26.0", "openai": "^6.33.0", "react": "^19.1.0", "react-dom": "^19.1.0", diff --git a/src/components/KGOnePanel.css b/src/components/KGOnePanel.css index d37a191..d7d79c9 100644 --- a/src/components/KGOnePanel.css +++ b/src/components/KGOnePanel.css @@ -62,6 +62,15 @@ border-bottom-color: #5a9fd4; } +.kgone-tab.is-disabled { + color: #666; + cursor: not-allowed; +} + +.kgone-tab.is-disabled:hover { + color: #666; +} + /* Scrollable body */ .kgone-panel-body { flex: 1; @@ -203,6 +212,64 @@ border-radius: 6px; } +.kgone-local-mode-card { + background-color: #252525; + border: 1px solid #3a3a3a; + border-radius: 6px; + padding: 10px; + display: flex; + flex-direction: column; + gap: 8px; +} + +.kgone-local-mode-title { + color: #e0e0e0; + font-size: 12px; + font-weight: 600; +} + +.kgone-local-mode-text { + color: #999; + font-size: 11px; + line-height: 1.45; +} + +.kgone-runtime-row { + display: flex; + flex-wrap: wrap; + gap: 6px; +} + +.kgone-provider-chip { + background-color: #313131; + border: 1px solid #444; + border-radius: 999px; + color: #bbb; + font-size: 10px; + padding: 4px 8px; +} + +.kgone-progress-block { + display: flex; + flex-direction: column; + gap: 6px; +} + +.kgone-progress-track { + width: 100%; + height: 8px; + background-color: #1d1d1d; + border: 1px solid #3a3a3a; + border-radius: 999px; + overflow: hidden; +} + +.kgone-progress-fill { + height: 100%; + background: linear-gradient(90deg, #5a9fd4 0%, #7cc2f1 100%); + transition: width 0.15s linear; +} + .kgone-region-info { background-color: #252525; border: 1px solid #3a3a3a; @@ -252,6 +319,32 @@ cursor: not-allowed; } +.kgone-btn-secondary { + background-color: #353535; + border: 1px solid #4a4a4a; + border-radius: 4px; + color: #ddd; + font-size: 11px; + font-weight: 600; + padding: 8px 10px; + cursor: pointer; + flex: 1; +} + +.kgone-btn-secondary:hover:not(:disabled) { + background-color: #404040; +} + +.kgone-btn-secondary:disabled { + color: #666; + cursor: not-allowed; +} + +.kgone-btn-danger { + border-color: #5a2a2a; + color: #f0b0b0; +} + /* ── Audio player ──────────────────────────────────────────────────────────── */ .kgone-audio-player { display: flex; @@ -404,4 +497,4 @@ font-size: 11px; cursor: pointer; user-select: none; -} \ No newline at end of file +} diff --git a/src/components/KGOnePanel.test.tsx b/src/components/KGOnePanel.test.tsx new file mode 100644 index 0000000..f63f5bf --- /dev/null +++ b/src/components/KGOnePanel.test.tsx @@ -0,0 +1,180 @@ +import React from 'react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import KGOnePanel from './KGOnePanel'; +import { KGAudioRegion } from '../core/region/KGAudioRegion'; +import { KGAudioTrack } from '../core/track/KGAudioTrack'; + +let kgoneEnabled = false; +let selectedRegionIds: string[] = []; +let localModelCached = false; +let localSeparationResult: Array<{ name: string; blob: Blob }> = []; + +const mockRefreshProjectState = vi.fn(); +const mockExecuteCommand = vi.fn(); + +vi.mock('../stores/projectStore', () => ({ + useProjectStore: () => ({ + selectedRegionIds, + projectName: 'Test Project', + bpm: 120, + keySignature: 'C major', + timeSignature: { numerator: 4, denominator: 4 }, + maxBars: 32, + refreshProjectState: mockRefreshProjectState, + }), +})); + +vi.mock('../core/config/ConfigManager', () => ({ + ConfigManager: { + instance: () => ({ + get: (key: string) => { + if (key === 'general.kgone.enabled') return kgoneEnabled; + if (key === 'general.kgone.base_url') return 'http://127.0.0.1:8000'; + return undefined; + }, + }), + }, +})); + +const audioRegion = new KGAudioRegion( + 'audio-region-1', + 'track-1', + 0, + 'Verse Stem', + 0, + 4, + 'audio-file-1', + 'verse.wav', + 2, + 0, +); +const audioTrack = new KGAudioTrack('Audio Track', 1); +audioTrack.setTrackIndex(0); +audioTrack.setRegions([audioRegion]); + +vi.mock('../core/KGCore', () => ({ + KGCore: { + instance: () => ({ + getCurrentProject: () => ({ + getTracks: () => [audioTrack], + getName: () => 'Test Project', + }), + executeCommand: mockExecuteCommand, + }), + }, +})); + +vi.mock('../core/io/KGAudioFileStorage', () => ({ + KGAudioFileStorage: { + loadAudioFile: vi.fn(async () => new ArrayBuffer(8)), + storeAudioFile: vi.fn(async () => undefined), + }, +})); + +vi.mock('../util/audioUtil', () => ({ + sliceAudioToWav: vi.fn(async (_buffer: ArrayBuffer) => _buffer), +})); + +vi.mock('../util/localSeparatorModelCache', () => ({ + LocalSeparatorModelCache: { + exists: vi.fn(async () => localModelCached), + download: vi.fn(async () => { + localModelCached = true; + }), + delete: vi.fn(async () => { + localModelCached = false; + }), + getArrayBuffer: vi.fn(async () => new ArrayBuffer(16)), + }, +})); + +vi.mock('../util/localSeparatorRuntime', () => ({ + detectLocalRuntimeSupport: () => ({ webgpuExposed: false }), + LocalOrtRuntimeManager: class { + constructor(private readonly options?: { onProviderChange?: (provider: string) => void }) {} + + reset() {} + + async ensureRuntime() { + this.options?.onProviderChange?.('cpu/wasm'); + return { provider: 'wasm', session: {} }; + } + }, +})); + +vi.mock('../util/localSeparatorRunner', () => ({ + runLocalSeparator: vi.fn(async ({ onProgress, onProviderChange }) => { + onProviderChange?.('cpu/wasm'); + onProgress({ stage: 'main', passLabel: 'Main pass', percent: 100, processedChunks: 1, totalChunks: 1 }); + return { + stems: localSeparationResult, + providerLabel: 'CPU/wasm', + debugSummary: {}, + }; + }), +})); + +describe('KGOnePanel local separator mode', () => { + beforeEach(() => { + kgoneEnabled = false; + selectedRegionIds = []; + localModelCached = false; + localSeparationResult = [ + { name: 'Instrumental', blob: new Blob(['instrumental'], { type: 'audio/wav' }) }, + { name: 'Vocals', blob: new Blob(['vocals'], { type: 'audio/wav' }) }, + ]; + mockRefreshProjectState.mockReset(); + mockExecuteCommand.mockReset(); + }); + + it('defaults to separator and shows other tabs as disabled in local mode', async () => { + render(); + + expect(await screen.findByText('Local Separator Mode')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Full Song' })).toBeDisabled(); + expect(screen.getByRole('button', { name: 'Remix' })).toBeDisabled(); + expect(screen.getByRole('button', { name: 'Repaint' })).toBeDisabled(); + expect(screen.getByRole('button', { name: 'Separator' })).not.toBeDisabled(); + expect(screen.getByRole('button', { name: 'Download Model' })).toBeInTheDocument(); + }); + + it('shows the single local separator model and advanced settings when the model is cached', async () => { + localModelCached = true; + selectedRegionIds = ['audio-region-1']; + + render(); + + await screen.findByText('Selected Region'); + const options = await screen.findAllByRole('option'); + expect(options).toHaveLength(1); + expect(options[0]).toHaveTextContent('Vocal and Instrument (Medium Accuracy)'); + + fireEvent.click(screen.getByRole('button', { name: /Advanced Settings/i })); + expect(screen.getByLabelText('Optional audio chunk duration (seconds)')).toBeInTheDocument(); + expect(screen.getByLabelText('MDX overlap')).toBeInTheDocument(); + }); + + it('prompts for an audio region when the model is cached but nothing is selected', async () => { + localModelCached = true; + + render(); + + expect(await screen.findByText(/Select an audio region on the timeline/)).toBeInTheDocument(); + }); + + it('renders local separation outputs after processing completes', async () => { + localModelCached = true; + selectedRegionIds = ['audio-region-1']; + + render(); + + fireEvent.click(await screen.findByRole('button', { name: 'Separate Stems' })); + + await waitFor(() => { + expect(screen.getByText('Instrumental')).toBeInTheDocument(); + expect(screen.getByText('Vocals')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument(); + }); + }); +}); diff --git a/src/components/KGOnePanel.tsx b/src/components/KGOnePanel.tsx index d53dbd2..26012dd 100644 --- a/src/components/KGOnePanel.tsx +++ b/src/components/KGOnePanel.tsx @@ -15,18 +15,31 @@ import type { KeySignature } from '../core/KGProject'; import { ImportStemsCommand } from '../core/commands'; import type { StemImportEntry } from '../core/commands'; import { showAlert } from '../util/dialogUtil'; +import { + LOCAL_SEPARATOR_MODEL_CONFIG, + LOCAL_SEPARATOR_MODEL_FILENAME, + LOCAL_SEPARATOR_MODEL_URL, +} from '../util/localSeparatorConfig'; +import { LocalSeparatorModelCache } from '../util/localSeparatorModelCache'; +import { runLocalSeparator } from '../util/localSeparatorRunner'; +import { LocalOrtRuntimeManager, detectLocalRuntimeSupport } from '../util/localSeparatorRuntime'; // ─── Types ──────────────────────────────────────────────────────────────────── type Tab = 'clip' | 'fullsong' | 'remix' | 'repaint' | 'separator'; +type KGOneMode = 'server' | 'local-separator'; type GenStatus = 'idle' | 'loading-model' | 'generating' | 'polling' | 'downloading' | 'done' | 'error'; -const SEPARATOR_MODELS = [ +const SERVER_SEPARATOR_MODELS = [ { label: 'Vocal and Instrument (Medium Accuracy)', value: 'UVR-MDX-NET-Inst_HQ_3.onnx' }, { label: 'Vocal and Instrument (High Accuracy)', value: 'MDX23C-8KFFT-InstVoc_HQ.ckpt' }, { label: 'Vocal, Drums, Bass, Guitar, Piano, and Others', value: 'htdemucs_6s.yaml' }, ] as const; +const LOCAL_SEPARATOR_MODELS = [ + { label: LOCAL_SEPARATOR_MODEL_CONFIG.displayName, value: LOCAL_SEPARATOR_MODEL_FILENAME }, +] as const; +const KGONE_TABS = ['fullsong', 'remix', 'repaint', 'separator'] as const; const CLIP_NOTES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']; @@ -59,6 +72,22 @@ function formatTime(sec: number): string { return `${m}:${s.toString().padStart(2, '0')}`; } +function formatKGOneTabLabel(tab: Tab): string { + if (tab === 'fullsong') return 'Full Song'; + if (tab === 'remix') return 'Remix'; + if (tab === 'repaint') return 'Repaint'; + return 'Separator'; +} + +export function getDefaultKGOneTab(mode: KGOneMode): Tab { + return mode === 'local-separator' ? 'separator' : 'fullsong'; +} + +export function getKGOneMode(): KGOneMode { + const enabled = (ConfigManager.instance().get('general.kgone.enabled') as boolean | undefined) ?? false; + return enabled ? 'server' : 'local-separator'; +} + // ─── Shared components ──────────────────────────────────────────────────────── interface ExpanderProps { @@ -830,18 +859,33 @@ function countRepaintTracks(sourceTrackName: string): number { return tracks.filter(t => pattern.test(t.getName())).length; } -const SeparatorTab: React.FC = () => { +const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => { const { selectedRegionIds, projectName, bpm, timeSignature, maxBars, refreshProjectState } = useProjectStore(); - const [model, setModel] = useState(SEPARATOR_MODELS[0].value); + const localOnlyMode = mode === 'local-separator'; + const availableSeparatorModels = localOnlyMode ? LOCAL_SEPARATOR_MODELS : SERVER_SEPARATOR_MODELS; + const [model, setModel] = useState(availableSeparatorModels[0].value); // Generation state const [genStatus, setGenStatus] = useState('idle'); const [genHint, setGenHint] = useState(''); const [errorMsg, setErrorMsg] = useState(''); const [stemAudioUrls, setStemAudioUrls] = useState>([]); + const runtimeSupport = useMemo(() => detectLocalRuntimeSupport(), []); + const [localProviderLabel, setLocalProviderLabel] = useState( + runtimeSupport.webgpuExposed ? 'webgpu available' : 'cpu/wasm only', + ); + const [isLocalModelCached, setIsLocalModelCached] = useState(false); + const [isCheckingLocalModel, setIsCheckingLocalModel] = useState(false); + const [isDownloadingLocalModel, setIsDownloadingLocalModel] = useState(false); + const [isDeletingLocalModel, setIsDeletingLocalModel] = useState(false); + const [localProgressPercent, setLocalProgressPercent] = useState(0); + const [localProgressText, setLocalProgressText] = useState(''); + const [localChunkDurationSeconds, setLocalChunkDurationSeconds] = useState(''); + const [localOverlap, setLocalOverlap] = useState(String(LOCAL_SEPARATOR_MODEL_CONFIG.defaults.overlap)); const abortRef = useRef(null); const taskIdRef = useRef(''); + const localRuntimeManagerRef = useRef(null); const originalRegionRef = useRef<{ regionName: string; trackName: string; @@ -861,6 +905,31 @@ const SeparatorTab: React.FC = () => { // eslint-disable-next-line react-hooks/exhaustive-deps }, []); + useEffect(() => { + setModel(availableSeparatorModels[0].value); + }, [availableSeparatorModels]); + + const refreshLocalModelCacheState = useCallback(async () => { + if (!localOnlyMode) return; + setIsCheckingLocalModel(true); + try { + setIsLocalModelCached(await LocalSeparatorModelCache.exists()); + } catch (err) { + console.error('[KGOne] Local model cache check failed:', err); + setErrorMsg(err instanceof Error ? err.message : String(err)); + } finally { + setIsCheckingLocalModel(false); + } + }, [localOnlyMode]); + + useEffect(() => { + if (!localOnlyMode) return; + localRuntimeManagerRef.current = new LocalOrtRuntimeManager({ + onProviderChange: provider => setLocalProviderLabel(provider), + }); + void refreshLocalModelCacheState(); + }, [localOnlyMode, refreshLocalModelCacheState]); + const selectedAudioRegion = useMemo(() => { if (!selectedRegionIds.length) return null; const project = KGCore.instance().getCurrentProject(); @@ -879,7 +948,57 @@ const SeparatorTab: React.FC = () => { const isGenerating = genStatus !== 'idle' && genStatus !== 'done' && genStatus !== 'error'; - const handleSeparate = useCallback(async () => { + const handleDownloadLocalModel = useCallback(async () => { + setIsDownloadingLocalModel(true); + setErrorMsg(''); + setLocalProgressPercent(0); + setLocalProgressText('Downloading local separator model...'); + + try { + await LocalSeparatorModelCache.download( + LOCAL_SEPARATOR_MODEL_URL, + LOCAL_SEPARATOR_MODEL_FILENAME, + progress => { + const receivedMb = (progress.receivedBytes / (1024 * 1024)).toFixed(1); + const totalMb = progress.totalBytes ? (progress.totalBytes / (1024 * 1024)).toFixed(1) : null; + setLocalProgressPercent(progress.totalBytes ? progress.percent : 0); + setLocalProgressText( + totalMb + ? `Downloading local separator model... ${receivedMb} / ${totalMb} MB` + : `Downloading local separator model... ${receivedMb} MB`, + ); + }, + ); + setLocalProgressPercent(100); + setLocalProgressText('Local separator model is ready.'); + await refreshLocalModelCacheState(); + } catch (err) { + setLocalProgressPercent(0); + setLocalProgressText(''); + setErrorMsg(err instanceof Error ? err.message : String(err)); + } finally { + setIsDownloadingLocalModel(false); + } + }, [refreshLocalModelCacheState]); + + const handleDeleteLocalModel = useCallback(async () => { + setIsDeletingLocalModel(true); + setErrorMsg(''); + try { + await LocalSeparatorModelCache.delete(); + localRuntimeManagerRef.current?.reset(); + setLocalProviderLabel(runtimeSupport.webgpuExposed ? 'webgpu available' : 'cpu/wasm only'); + setLocalProgressPercent(0); + setLocalProgressText(''); + await refreshLocalModelCacheState(); + } catch (err) { + setErrorMsg(err instanceof Error ? err.message : String(err)); + } finally { + setIsDeletingLocalModel(false); + } + }, [refreshLocalModelCacheState, runtimeSupport.webgpuExposed]); + + const handleSeparateServer = useCallback(async () => { if (!selectedAudioRegion) return; // Capture snapshot before anything changes — selection may shift during generation @@ -1052,7 +1171,107 @@ const SeparatorTab: React.FC = () => { setGenStatus('error'); setGenHint(''); } - }, [selectedAudioRegion, projectName, model, stemAudioUrls]); + }, [selectedAudioRegion, projectName, model, stemAudioUrls, bpm]); + + const handleSeparateLocal = useCallback(async () => { + if (!selectedAudioRegion || !isLocalModelCached) return; + + originalRegionRef.current = { + regionName: selectedAudioRegion.region.getName(), + trackName: selectedAudioRegion.trackName, + startFromBeat: selectedAudioRegion.region.getStartFromBeat(), + trackIndex: selectedAudioRegion.trackIndex, + }; + setImportError(''); + stemAudioUrls.forEach(s => URL.revokeObjectURL(s.url)); + setStemAudioUrls([]); + setErrorMsg(''); + setGenStatus('loading-model'); + setLocalProgressPercent(0); + setLocalProgressText('Preparing ONNX Runtime session...'); + setLocalProviderLabel(runtimeSupport.webgpuExposed ? 'webgpu available' : 'cpu/wasm only'); + + try { + const modelBuffer = await LocalSeparatorModelCache.getArrayBuffer(); + const runtimeManager = localRuntimeManagerRef.current ?? new LocalOrtRuntimeManager({ + onProviderChange: provider => setLocalProviderLabel(provider), + }); + localRuntimeManagerRef.current = runtimeManager; + const runtime = await runtimeManager.ensureRuntime(LOCAL_SEPARATOR_MODEL_CONFIG, new Uint8Array(modelBuffer)); + + setGenStatus('generating'); + setLocalProgressPercent(3); + setLocalProgressText('Reading audio file...'); + + const audioFileId = selectedAudioRegion.region.getAudioFileId(); + const clipStart = selectedAudioRegion.region.getClipStartOffsetSeconds(); + const fullDuration = selectedAudioRegion.region.getAudioDurationSeconds(); + const regionLengthSec = selectedAudioRegion.region.getLength() * (60 / bpm); + const effectiveDuration = Math.min(regionLengthSec, fullDuration - clipStart); + + const rawBuffer = await KGAudioFileStorage.loadAudioFile(projectName, audioFileId); + const needsSlice = clipStart > 0.01 || effectiveDuration < fullDuration - 0.01; + const inputBuffer = needsSlice + ? await sliceAudioToWav(rawBuffer, clipStart, effectiveDuration) + : rawBuffer; + + setGenStatus('polling'); + setLocalProgressPercent(5); + setLocalProgressText('Running browser separation...'); + + const chunkDuration = localChunkDurationSeconds.trim() + ? Number.parseFloat(localChunkDurationSeconds) + : null; + const overlapValue = Number.parseFloat(localOverlap); + + const result = await runLocalSeparator({ + session: runtime.session, + runtimeProvider: runtime.provider, + modelConfig: LOCAL_SEPARATOR_MODEL_CONFIG, + audioBuffer: inputBuffer, + chunkDurationSeconds: Number.isFinite(chunkDuration) && (chunkDuration ?? 0) > 0 ? chunkDuration : null, + overlap: Number.isFinite(overlapValue) ? overlapValue : LOCAL_SEPARATOR_MODEL_CONFIG.defaults.overlap, + onProviderChange: provider => setLocalProviderLabel(provider), + onProgress: progress => { + setLocalProgressPercent(progress.percent); + const chunkSuffix = progress.totalChunks ? ` (${progress.processedChunks}/${progress.totalChunks} chunks)` : ''; + setLocalProgressText(`${progress.passLabel}${chunkSuffix}`); + }, + }); + + taskIdRef.current = `local_${Date.now()}`; + const nextStemAudioUrls = result.stems.map(stem => ({ + name: stem.name, + url: URL.createObjectURL(stem.blob), + })); + setStemAudioUrls(nextStemAudioUrls); + setGenStatus('done'); + setLocalProgressPercent(100); + setLocalProgressText('Separation complete.'); + kgoneLog('RES', 'local separator summary', result.debugSummary); + } catch (err) { + console.error('[KGOne] Local separator error:', err); + setGenStatus('error'); + setErrorMsg(err instanceof Error ? err.message : String(err)); + } + }, [ + selectedAudioRegion, + isLocalModelCached, + stemAudioUrls, + runtimeSupport.webgpuExposed, + bpm, + projectName, + localChunkDurationSeconds, + localOverlap, + ]); + + const handleSeparate = useCallback(async () => { + if (localOnlyMode) { + await handleSeparateLocal(); + return; + } + await handleSeparateServer(); + }, [handleSeparateLocal, handleSeparateServer, localOnlyMode]); const handleImportAll = useCallback(async () => { const snap = originalRegionRef.current; @@ -1068,9 +1287,9 @@ const SeparatorTab: React.FC = () => { const stems: StemImportEntry[] = await Promise.all( stemAudioUrls.map(async (stem) => { const blob = await fetch(stem.url).then(r => r.blob()); - const fileName = `KGOne_Stem_${stem.name}_${taskIdRef.current}.mp3`; + const fileName = `KGOne_Stem_${stem.name}_${taskIdRef.current}.${localOnlyMode ? 'wav' : 'mp3'}`; const fileId = `kgone_stem_${Date.now()}_${Math.random().toString(36).slice(2, 7)}`; - const audioFile = new File([blob], fileName, { type: 'audio/mpeg' }); + const audioFile = new File([blob], fileName, { type: localOnlyMode ? 'audio/wav' : 'audio/mpeg' }); const arrayBuffer = await blob.arrayBuffer(); const toneBuffer = new Tone.ToneAudioBuffer(); @@ -1114,9 +1333,19 @@ const SeparatorTab: React.FC = () => { } finally { setIsImporting(false); } - }, [stemAudioUrls, projectName, maxBars, refreshProjectState]); + }, [stemAudioUrls, projectName, maxBars, refreshProjectState, localOnlyMode]); const btnLabel = () => { + if (localOnlyMode) { + switch (genStatus) { + case 'loading-model': return 'Preparing local model...'; + case 'generating': return 'Preparing audio...'; + case 'polling': return 'Separating locally...'; + case 'downloading': return 'Finalizing...'; + default: return 'Separate Stems'; + } + } + switch (genStatus) { case 'loading-model': return 'Loading model...'; case 'generating': return 'Preparing upload...'; @@ -1128,6 +1357,68 @@ const SeparatorTab: React.FC = () => { return ( <> + {localOnlyMode && ( +
+
Local Separator Mode
+
+ Only Vocal and Instrument (Medium Accuracy) is available while not integrated with K.G.One Music Studio server. + Processing in local may take long time depending on your hardware. When fallback to CPU happens, the webpage may + temporarily hang with little or no UI response until processing advances. +
+
+
Provider: {localProviderLabel}
+
Model: {isLocalModelCached ? 'downloaded' : 'not downloaded'}
+
+ {(localProgressText || isCheckingLocalModel) && ( +
+
+
+
+
+ {isCheckingLocalModel ? 'Checking local model cache...' : localProgressText} +
+
+ )} +
+ {!isLocalModelCached ? ( + + ) : ( + <> + + + + )} +
+
+ )} + {selectedAudioRegion ? ( <>
@@ -1139,14 +1430,44 @@ const SeparatorTab: React.FC = () => {
- setModel(e.target.value as typeof SERVER_SEPARATOR_MODELS[number]['value'])}> + {availableSeparatorModels.map(m => ( ))}
+ {localOnlyMode && ( + +
+ + setLocalChunkDurationSeconds(e.target.value)} + placeholder="Leave blank to process the full region" + /> +
+
+ + setLocalOverlap(e.target.value)} + /> +
+
+ )} + {/* Stem audio players — shown once separation is complete */} {stemAudioUrls.length > 0 && (
@@ -1156,7 +1477,7 @@ const SeparatorTab: React.FC = () => {
@@ -1195,7 +1516,7 @@ const SeparatorTab: React.FC = () => { {/* Status hint below button */} - {genHint &&
{genHint}
} + {(localOnlyMode ? localProgressText : genHint) && ( +
{localOnlyMode ? localProgressText : genHint}
+ )} ) : (
- Select an audio region on the timeline to extract stems from it. - Only audio regions are supported — MIDI regions cannot be separated. + {localOnlyMode && !isLocalModelCached + ? 'Download the local separator model, then select an audio region on the timeline to extract stems from it.' + : 'Select an audio region on the timeline to extract stems from it. Only audio regions are supported — MIDI regions cannot be separated.'}
)} -
- Powered by UVR5 CLI -
+ {!localOnlyMode && ( +
+ Powered by UVR5 CLI +
+ )} ); }; @@ -2228,8 +2554,21 @@ interface KGOnePanelProps { } const KGOnePanel: React.FC = ({ isVisible }) => { - const [activeTab, setActiveTab] = useState('fullsong'); + const mode = getKGOneMode(); + const [activeTab, setActiveTab] = useState(getDefaultKGOneTab(mode)); const { bpm, keySignature } = useProjectStore(); + const disabledTabs = mode === 'local-separator' + ? new Set(['fullsong', 'remix', 'repaint']) + : new Set(); + + useEffect(() => { + setActiveTab(current => { + if (disabledTabs.has(current)) { + return 'separator'; + } + return getDefaultKGOneTab(mode) === 'separator' ? 'separator' : current; + }); + }, [mode]); return (
@@ -2239,16 +2578,17 @@ const KGOnePanel: React.FC = ({ isVisible }) => {
{/* Clip tab temporarily disabled, will enable in the future */} - {(['fullsong', 'remix', 'repaint', 'separator'] as const).map(tab => ( + {KGONE_TABS.map(tab => ( ))}
@@ -2259,7 +2599,7 @@ const KGOnePanel: React.FC = ({ isVisible }) => { {activeTab === 'fullsong' && } {activeTab === 'remix' && } {activeTab === 'repaint' && } - {activeTab === 'separator' && } + {activeTab === 'separator' && }
); diff --git a/src/components/Toolbar.tsx b/src/components/Toolbar.tsx index af37e84..b4279be 100644 --- a/src/components/Toolbar.tsx +++ b/src/components/Toolbar.tsx @@ -961,8 +961,6 @@ const Toolbar: React.FC = () => { }; // K.G.One panel toggle - const isKGOneEnabled = ConfigManager.instance().get('general.kgone.enabled') as boolean ?? false; - const handleKGOneClick = () => { if (DEBUG_MODE.TOOLBAR) { console.log("K.G.One button clicked"); @@ -1241,11 +1239,9 @@ const Toolbar: React.FC = () => {
diff --git a/src/core/io/LocalSeparatorModelCache.test.ts b/src/core/io/LocalSeparatorModelCache.test.ts new file mode 100644 index 0000000..0642e25 --- /dev/null +++ b/src/core/io/LocalSeparatorModelCache.test.ts @@ -0,0 +1,151 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { LocalSeparatorModelCache } from '../../util/localSeparatorModelCache'; + +class MockWritableFileStream { + private readonly handle: MockFileSystemFileHandle; + private chunks: Uint8Array[] = []; + + constructor(handle: MockFileSystemFileHandle) { + this.handle = handle; + } + + async write(content: ArrayBuffer | ArrayBufferView): Promise { + const bytes = content instanceof ArrayBuffer + ? new Uint8Array(content) + : new Uint8Array(content.buffer, content.byteOffset, content.byteLength); + this.chunks.push(new Uint8Array(bytes)); + } + + async close(): Promise { + const total = this.chunks.reduce((sum, chunk) => sum + chunk.byteLength, 0); + const merged = new Uint8Array(total); + let offset = 0; + for (const chunk of this.chunks) { + merged.set(chunk, offset); + offset += chunk.byteLength; + } + this.handle.setContent(merged); + } + + async abort(): Promise { + this.chunks = []; + } +} + +class MockFileSystemFileHandle { + kind = 'file' as const; + private content = new Uint8Array(); + + constructor(public readonly name: string) {} + + setContent(content: Uint8Array): void { + this.content = content; + } + + async getFile(): Promise { + return { + arrayBuffer: async () => this.content.buffer.slice(0), + } as unknown as File; + } + + async createWritable(): Promise { + return new MockWritableFileStream(this); + } +} + +class MockFileSystemDirectoryHandle { + kind = 'directory' as const; + private entries = new Map(); + + constructor(public readonly name: string) {} + + async getDirectoryHandle(name: string, options?: { create?: boolean }): Promise { + let entry = this.entries.get(name); + if (!entry || entry.kind !== 'directory') { + if (!options?.create) { + throw new DOMException(`Directory "${name}" not found`, 'NotFoundError'); + } + entry = new MockFileSystemDirectoryHandle(name); + this.entries.set(name, entry); + } + return entry as MockFileSystemDirectoryHandle; + } + + async getFileHandle(name: string, options?: { create?: boolean }): Promise { + let entry = this.entries.get(name); + if (!entry || entry.kind !== 'file') { + if (!options?.create) { + throw new DOMException(`File "${name}" not found`, 'NotFoundError'); + } + entry = new MockFileSystemFileHandle(name); + this.entries.set(name, entry); + } + return entry as MockFileSystemFileHandle; + } + + async removeEntry(name: string): Promise { + if (!this.entries.has(name)) { + throw new DOMException(`Entry "${name}" not found`, 'NotFoundError'); + } + this.entries.delete(name); + } + + clear(): void { + this.entries.clear(); + } +} + +const mockRoot = new MockFileSystemDirectoryHandle('root'); + +vi.stubGlobal('navigator', { + ...navigator, + storage: { + getDirectory: vi.fn(async () => mockRoot), + }, +}); + +describe('LocalSeparatorModelCache', () => { + beforeEach(() => { + mockRoot.clear(); + vi.restoreAllMocks(); + }); + + it('downloads and stores a model in OPFS cache', async () => { + vi.stubGlobal('fetch', vi.fn(async () => new Response(new Uint8Array([1, 2, 3]), { + status: 200, + headers: { 'Content-Length': '3' }, + }))); + + await LocalSeparatorModelCache.download('https://example.com/model.onnx', 'model.onnx'); + + expect(await LocalSeparatorModelCache.exists('model.onnx')).toBe(true); + const buffer = await LocalSeparatorModelCache.getArrayBuffer('model.onnx'); + expect(Array.from(new Uint8Array(buffer))).toEqual([1, 2, 3]); + }); + + it('replaces a broken cached file on redownload', async () => { + vi.stubGlobal('fetch', vi.fn(async () => new Response(new Uint8Array([1]), { + status: 200, + headers: { 'Content-Length': '1' }, + }))); + await LocalSeparatorModelCache.download('https://example.com/model.onnx', 'model.onnx'); + + vi.stubGlobal('fetch', vi.fn(async () => new Response(new Uint8Array([9, 8, 7, 6]), { + status: 200, + headers: { 'Content-Length': '4' }, + }))); + await LocalSeparatorModelCache.download('https://example.com/model.onnx', 'model.onnx'); + + const buffer = await LocalSeparatorModelCache.getArrayBuffer('model.onnx'); + expect(Array.from(new Uint8Array(buffer))).toEqual([9, 8, 7, 6]); + }); + + it('deletes the cached model file', async () => { + vi.stubGlobal('fetch', vi.fn(async () => new Response(new Uint8Array([1, 2]), { status: 200 }))); + await LocalSeparatorModelCache.download('https://example.com/model.onnx', 'model.onnx'); + + await LocalSeparatorModelCache.delete('model.onnx'); + + expect(await LocalSeparatorModelCache.exists('model.onnx')).toBe(false); + }); +}); diff --git a/src/util/localSeparatorConfig.ts b/src/util/localSeparatorConfig.ts new file mode 100644 index 0000000..752f468 --- /dev/null +++ b/src/util/localSeparatorConfig.ts @@ -0,0 +1,31 @@ +import type { LocalSeparatorModelConfig } from './localSeparatorTypes'; + +export const LOCAL_SEPARATOR_MODEL_URL = + 'https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx'; + +export const LOCAL_SEPARATOR_MODEL_FILENAME = 'UVR-MDX-NET-Inst_HQ_3.onnx'; + +export const LOCAL_SEPARATOR_MODEL_CONFIG: LocalSeparatorModelConfig = { + filename: LOCAL_SEPARATOR_MODEL_FILENAME, + displayName: 'Vocal and Instrument (Medium Accuracy)', + status: 'ready', + defaults: { + sampleRate: 44100, + hopLength: 1024, + segmentSize: 256, + overlap: 0.25, + batchSize: 1, + enableDenoise: false, + invertUsingSpec: false, + normalizationThreshold: 0.9, + amplificationThreshold: 0, + matchMixOverlap: 0.02, + }, + metadata: { + compensate: 1.021, + mdx_dim_f_set: 3072, + mdx_dim_t_set: 8, + mdx_n_fft_scale_set: 7680, + primary_stem: 'Instrumental', + }, +}; diff --git a/src/util/localSeparatorCpuDsp.ts b/src/util/localSeparatorCpuDsp.ts new file mode 100644 index 0000000..cda0ac3 --- /dev/null +++ b/src/util/localSeparatorCpuDsp.ts @@ -0,0 +1,139 @@ +import type { LocalSeparatorModelConfig, StereoChannels } from './localSeparatorTypes'; +import { FFT, createWindowCache, getHannPeriodic, index4d, reflectPad } from './localSeparatorShared'; + +interface SpectrogramPayload { + data: Float32Array; + dims: number[]; + frames?: number; +} + +export class LocalSeparatorCpuDsp { + public readonly window: Float32Array; + public readonly dimF: number; + public readonly forwardFft: FFT; + + private readonly nFft: number; + private readonly hopLength: number; + private readonly trim: number; + private readonly numFreqBins: number; + private readonly inverseFft: FFT; + + constructor(config: LocalSeparatorModelConfig) { + this.nFft = config.metadata.mdx_n_fft_scale_set; + this.hopLength = config.defaults.hopLength; + this.dimF = config.metadata.mdx_dim_f_set; + this.trim = Math.floor(this.nFft / 2); + this.numFreqBins = Math.floor(this.nFft / 2) + 1; + const windowCache = createWindowCache(); + this.window = getHannPeriodic(this.nFft, windowCache); + this.forwardFft = new FFT(this.nFft); + this.inverseFft = new FFT(this.nFft); + } + + public async forwardStereo(leftChunk: Float32Array, rightChunk: Float32Array): Promise { + const paddedLeft = reflectPad(leftChunk, this.trim, this.trim); + const paddedRight = reflectPad(rightChunk, this.trim, this.trim); + const frames = Math.floor((paddedLeft.length - this.nFft) / this.hopLength) + 1; + const tensor = new Float32Array(4 * this.dimF * frames); + const dims = [1, 4, this.dimF, frames]; + + for (let frameIndex = 0; frameIndex < frames; frameIndex += 1) { + const offset = frameIndex * this.hopLength; + const leftSpectrum = this.frameSpectrum(paddedLeft, offset); + const rightSpectrum = this.frameSpectrum(paddedRight, offset); + + for (let freq = 0; freq < this.dimF; freq += 1) { + tensor[index4d(dims, 0, 0, freq, frameIndex)] = freq < 3 ? 0 : leftSpectrum.real[freq]; + tensor[index4d(dims, 0, 1, freq, frameIndex)] = freq < 3 ? 0 : leftSpectrum.imag[freq]; + tensor[index4d(dims, 0, 2, freq, frameIndex)] = freq < 3 ? 0 : rightSpectrum.real[freq]; + tensor[index4d(dims, 0, 3, freq, frameIndex)] = freq < 3 ? 0 : rightSpectrum.imag[freq]; + } + } + + return { data: tensor, dims, frames }; + } + + public async inverseStereo(spectrogramPayload: SpectrogramPayload): Promise { + const spectrogram = spectrogramPayload.data; + const dims = spectrogramPayload.dims; + const [, channels, freqBins, frames] = dims; + if (channels !== 4) { + throw new Error(`Expected 4 channels in MDX spectrogram, got ${channels}`); + } + + const outputLength = ((frames - 1) * this.hopLength) + this.nFft; + const left = new Float64Array(outputLength); + const right = new Float64Array(outputLength); + const leftWindowSums = new Float64Array(outputLength); + const rightWindowSums = new Float64Array(outputLength); + + for (let frameIndex = 0; frameIndex < frames; frameIndex += 1) { + const leftFrame = this.istftFrame(spectrogram, dims, 0, 1, frameIndex, freqBins); + const rightFrame = this.istftFrame(spectrogram, dims, 2, 3, frameIndex, freqBins); + const frameOffset = frameIndex * this.hopLength; + + for (let i = 0; i < this.nFft; i += 1) { + const weightedLeft = leftFrame[i] * this.window[i]; + const weightedRight = rightFrame[i] * this.window[i]; + left[frameOffset + i] += weightedLeft; + right[frameOffset + i] += weightedRight; + const weight = this.window[i] * this.window[i]; + leftWindowSums[frameOffset + i] += weight; + rightWindowSums[frameOffset + i] += weight; + } + } + + const normalizedLeft = new Float32Array(outputLength - (this.trim * 2)); + const normalizedRight = new Float32Array(outputLength - (this.trim * 2)); + for (let i = this.trim; i < outputLength - this.trim; i += 1) { + const outIndex = i - this.trim; + normalizedLeft[outIndex] = leftWindowSums[i] > 1e-8 ? left[i] / leftWindowSums[i] : 0; + normalizedRight[outIndex] = rightWindowSums[i] > 1e-8 ? right[i] / rightWindowSums[i] : 0; + } + + return [normalizedLeft, normalizedRight]; + } + + public dispose(): void {} + + private frameSpectrum(signal: Float32Array, offset: number): { real: Float64Array; imag: Float64Array } { + const real = new Float64Array(this.nFft); + const imag = new Float64Array(this.nFft); + for (let i = 0; i < this.nFft; i += 1) { + real[i] = signal[offset + i] * this.window[i]; + } + this.forwardFft.transform(real, imag); + return { real, imag }; + } + + private istftFrame( + spectrogram: Float32Array, + dims: number[], + realChannel: number, + imagChannel: number, + frameIndex: number, + freqBins: number, + ): Float32Array { + const real = new Float64Array(this.nFft); + const imag = new Float64Array(this.nFft); + + for (let freq = 0; freq < freqBins; freq += 1) { + real[freq] = spectrogram[index4d(dims, 0, realChannel, freq, frameIndex)]; + imag[freq] = spectrogram[index4d(dims, 0, imagChannel, freq, frameIndex)]; + } + + for (let freq = 1; freq < this.numFreqBins - 1; freq += 1) { + const mirrored = this.nFft - freq; + real[mirrored] = real[freq]; + imag[mirrored] = -imag[freq]; + } + + this.inverseFft.inverse(real, imag); + + const frame = new Float32Array(this.nFft); + for (let i = 0; i < this.nFft; i += 1) { + frame[i] = real[i]; + } + return frame; + } +} diff --git a/src/util/localSeparatorGpuDsp.ts b/src/util/localSeparatorGpuDsp.ts new file mode 100644 index 0000000..dad4ba0 --- /dev/null +++ b/src/util/localSeparatorGpuDsp.ts @@ -0,0 +1,220 @@ +import type { LocalSeparatorModelConfig } from './localSeparatorTypes'; +import { LocalSeparatorCpuDsp } from './localSeparatorCpuDsp'; +import { reflectPad } from './localSeparatorShared'; + +type GPUDeviceLike = any; +type GPUBufferLike = any; +type GPUComputePipelineLike = any; + +declare const GPUBufferUsage: any; +declare const GPUMapMode: any; + +const FRAMING_SHADER = ` +struct Params { + nfft: u32, + hop: u32, + frames: u32, + paddedLength: u32, +} + +@group(0) @binding(0) var leftInput: array; +@group(0) @binding(1) var rightInput: array; +@group(0) @binding(2) var window: array; +@group(0) @binding(3) var output: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(256) +fn main(@builtin(global_invocation_id) gid: vec3) { + let index = gid.x; + let total = params.frames * params.nfft * 2u; + if (index >= total) { + return; + } + + let sample = index % params.nfft; + let frame = (index / params.nfft) % params.frames; + let channel = index / (params.nfft * params.frames); + let sourceIndex = frame * params.hop + sample; + let sampleValue = select(leftInput[sourceIndex], rightInput[sourceIndex], channel == 1u); + output[index] = sampleValue * window[sample]; +} +`; + +function alignTo(value: number, alignment: number): number { + return Math.ceil(value / alignment) * alignment; +} + +async function readBuffer(device: GPUDeviceLike, sourceBuffer: GPUBufferLike, size: number): Promise { + const readBuffer = device.createBuffer({ + size: alignTo(size, 4), + usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, + }); + + const encoder = device.createCommandEncoder(); + encoder.copyBufferToBuffer(sourceBuffer, 0, readBuffer, 0, size); + device.queue.submit([encoder.finish()]); + + await readBuffer.mapAsync(GPUMapMode.READ); + const copy = new Float32Array(readBuffer.getMappedRange().slice(0)); + readBuffer.unmap(); + readBuffer.destroy(); + return copy; +} + +export class LocalSeparatorGpuDsp { + private readonly device: GPUDeviceLike; + private readonly cpuDsp: LocalSeparatorCpuDsp; + private readonly nFft: number; + private readonly hopLength: number; + private readonly trim: number; + private windowBuffer: GPUBufferLike | null = null; + private paramBuffer: GPUBufferLike | null = null; + private readonly pipeline: GPUComputePipelineLike; + + public static async create(config: LocalSeparatorModelConfig): Promise { + if (!('gpu' in navigator)) { + throw new Error('WebGPU is not available for GPU DSP.'); + } + + const adapter = await (navigator as { gpu?: { requestAdapter: (options: { powerPreference: string }) => Promise } }).gpu?.requestAdapter({ + powerPreference: 'high-performance', + }); + if (!adapter) { + throw new Error('No WebGPU adapter was available for GPU DSP.'); + } + + const device = await adapter.requestDevice(); + return new LocalSeparatorGpuDsp(config, device); + } + + private constructor(config: LocalSeparatorModelConfig, device: GPUDeviceLike) { + this.device = device; + this.cpuDsp = new LocalSeparatorCpuDsp(config); + this.nFft = config.metadata.mdx_n_fft_scale_set; + this.hopLength = config.defaults.hopLength; + this.trim = Math.floor(this.nFft / 2); + this.pipeline = device.createComputePipeline({ + layout: 'auto', + compute: { + module: device.createShaderModule({ code: FRAMING_SHADER }), + entryPoint: 'main', + }, + }); + } + + public async forwardStereo(leftChunk: Float32Array, rightChunk: Float32Array): Promise<{ + data: Float32Array; + dims: number[]; + frames: number; + }> { + this.ensureStaticBuffers(); + + const paddedLeft = reflectPad(leftChunk, this.trim, this.trim); + const paddedRight = reflectPad(rightChunk, this.trim, this.trim); + const frames = Math.floor((paddedLeft.length - this.nFft) / this.hopLength) + 1; + + const leftBuffer = this.device.createBuffer({ + size: paddedLeft.byteLength, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + }); + const rightBuffer = this.device.createBuffer({ + size: paddedRight.byteLength, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + }); + const framedSize = frames * this.nFft * 2 * 4; + const outputBuffer = this.device.createBuffer({ + size: framedSize, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC, + }); + + this.device.queue.writeBuffer(leftBuffer, 0, paddedLeft); + this.device.queue.writeBuffer(rightBuffer, 0, paddedRight); + this.device.queue.writeBuffer(this.paramBuffer!, 0, new Uint32Array([this.nFft, this.hopLength, frames, paddedLeft.length])); + + const bindGroup = this.device.createBindGroup({ + layout: this.pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: leftBuffer } }, + { binding: 1, resource: { buffer: rightBuffer } }, + { binding: 2, resource: { buffer: this.windowBuffer! } }, + { binding: 3, resource: { buffer: outputBuffer } }, + { binding: 4, resource: { buffer: this.paramBuffer! } }, + ], + }); + + const encoder = this.device.createCommandEncoder(); + const pass = encoder.beginComputePass(); + pass.setPipeline(this.pipeline); + pass.setBindGroup(0, bindGroup); + pass.dispatchWorkgroups(Math.ceil((frames * this.nFft * 2) / 256)); + pass.end(); + this.device.queue.submit([encoder.finish()]); + + const framed = await readBuffer(this.device, outputBuffer, framedSize); + + leftBuffer.destroy(); + rightBuffer.destroy(); + outputBuffer.destroy(); + + return this.packFramedAudio(framed, frames); + } + + public async inverseStereo(payload: { data: Float32Array; dims: number[]; frames?: number }) { + return this.cpuDsp.inverseStereo(payload); + } + + public dispose(): void { + this.cpuDsp.dispose(); + this.windowBuffer?.destroy(); + this.paramBuffer?.destroy(); + } + + private ensureStaticBuffers(): void { + if (!this.windowBuffer) { + const window = this.cpuDsp.window; + this.windowBuffer = this.device.createBuffer({ + size: window.byteLength, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + }); + this.device.queue.writeBuffer(this.windowBuffer, 0, window); + } + + if (!this.paramBuffer) { + this.paramBuffer = this.device.createBuffer({ + size: 16, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + }); + } + } + + private packFramedAudio(framed: Float32Array, frames: number): { data: Float32Array; dims: number[]; frames: number } { + const tensor = new Float32Array(4 * this.cpuDsp.dimF * frames); + const dims = [1, 4, this.cpuDsp.dimF, frames]; + + for (let frameIndex = 0; frameIndex < frames; frameIndex += 1) { + const leftOffset = frameIndex * this.nFft; + const rightOffset = (frames * this.nFft) + leftOffset; + const leftSpectrum = this.fftFrame(framed, leftOffset); + const rightSpectrum = this.fftFrame(framed, rightOffset); + + for (let freq = 0; freq < this.cpuDsp.dimF; freq += 1) { + tensor[((freq * frames) + frameIndex)] = freq < 3 ? 0 : leftSpectrum.real[freq]; + tensor[(this.cpuDsp.dimF * frames) + ((freq * frames) + frameIndex)] = freq < 3 ? 0 : leftSpectrum.imag[freq]; + tensor[(2 * this.cpuDsp.dimF * frames) + ((freq * frames) + frameIndex)] = freq < 3 ? 0 : rightSpectrum.real[freq]; + tensor[(3 * this.cpuDsp.dimF * frames) + ((freq * frames) + frameIndex)] = freq < 3 ? 0 : rightSpectrum.imag[freq]; + } + } + + return { data: tensor, dims, frames }; + } + + private fftFrame(framed: Float32Array, offset: number): { real: Float64Array; imag: Float64Array } { + const real = new Float64Array(this.nFft); + const imag = new Float64Array(this.nFft); + for (let i = 0; i < this.nFft; i += 1) { + real[i] = framed[offset + i]; + } + this.cpuDsp.forwardFft.transform(real, imag); + return { real, imag }; + } +} diff --git a/src/util/localSeparatorModelCache.ts b/src/util/localSeparatorModelCache.ts new file mode 100644 index 0000000..a4990ac --- /dev/null +++ b/src/util/localSeparatorModelCache.ts @@ -0,0 +1,126 @@ +import { LOCAL_SEPARATOR_MODEL_FILENAME } from './localSeparatorConfig'; + +export interface ModelDownloadProgress { + receivedBytes: number; + totalBytes: number | null; + percent: number; +} + +export class LocalSeparatorModelCache { + private static readonly MODELS_DIR = 'models'; + private static readonly TEMP_SUFFIX = '.download'; + + public static async exists(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { + try { + const modelsDir = await this.getModelsDir(); + await modelsDir.getFileHandle(filename); + return true; + } catch { + return false; + } + } + + public static async getFile(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { + const modelsDir = await this.getModelsDir(); + const fileHandle = await modelsDir.getFileHandle(filename); + return fileHandle.getFile(); + } + + public static async getArrayBuffer(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { + const file = await this.getFile(filename); + return file.arrayBuffer(); + } + + public static async delete(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { + try { + const modelsDir = await this.getModelsDir(); + await modelsDir.removeEntry(filename); + } catch { + // Ignore missing file cleanup. + } + + try { + const modelsDir = await this.getModelsDir(); + await modelsDir.removeEntry(`${filename}${this.TEMP_SUFFIX}`); + } catch { + // Ignore missing temp file cleanup. + } + } + + public static async download( + sourceUrl: string, + filename: string = LOCAL_SEPARATOR_MODEL_FILENAME, + onProgress?: (progress: ModelDownloadProgress) => void, + ): Promise { + const response = await fetch(sourceUrl); + if (!response.ok) { + throw new Error(`Model download failed (${response.status})`); + } + + const modelsDir = await this.getModelsDir(); + const tempName = `${filename}${this.TEMP_SUFFIX}`; + await this.delete(filename); + + const tempHandle = await modelsDir.getFileHandle(tempName, { create: true }); + const writable = await tempHandle.createWritable(); + + try { + const totalBytesHeader = response.headers.get('Content-Length'); + const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : null; + + if (!response.body) { + const buffer = await response.arrayBuffer(); + await writable.write(buffer); + onProgress?.({ + receivedBytes: buffer.byteLength, + totalBytes, + percent: 100, + }); + } else { + const reader = response.body.getReader(); + let receivedBytes = 0; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + await writable.write(value); + receivedBytes += value.byteLength; + onProgress?.({ + receivedBytes, + totalBytes, + percent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0, + }); + } + } + } catch (error) { + await writable.abort(); + await this.delete(filename); + throw error; + } + + await writable.close(); + + const finalHandle = await modelsDir.getFileHandle(filename, { create: true }); + const finalWritable = await finalHandle.createWritable(); + try { + const tempFile = await tempHandle.getFile(); + await finalWritable.write(await tempFile.arrayBuffer()); + await finalWritable.close(); + } catch (error) { + await finalWritable.abort(); + throw error; + } finally { + try { + await modelsDir.removeEntry(tempName); + } catch { + // Ignore temp cleanup errors. + } + } + } + + private static async getModelsDir(): Promise { + const root = await navigator.storage.getDirectory(); + return root.getDirectoryHandle(this.MODELS_DIR, { create: true }); + } +} diff --git a/src/util/localSeparatorRunner.ts b/src/util/localSeparatorRunner.ts new file mode 100644 index 0000000..8419093 --- /dev/null +++ b/src/util/localSeparatorRunner.ts @@ -0,0 +1,589 @@ +import * as ort from 'onnxruntime-web/webgpu'; +import { LocalSeparatorCpuDsp } from './localSeparatorCpuDsp'; +import { LocalSeparatorGpuDsp } from './localSeparatorGpuDsp'; +import { LocalSeparatorTimingCollector } from './localSeparatorTiming'; +import type { + LocalRuntimeProvider, + LocalSeparatorModelConfig, + LocalSeparatorProgress, + StereoChannels, +} from './localSeparatorTypes'; +import { + concatFloat32, + createWindowCache, + getHanning, + negateArray, + normalizeChannels, + scaleChannels, + sliceChannels, +} from './localSeparatorShared'; + +const SAMPLE_RATE = 44100; + +function localSeparatorLog(message: string, payload?: unknown): void { + if (payload === undefined) { + console.log(`[localSeparator] ${message}`); + return; + } + console.log(`[localSeparator] ${message}`, payload); +} + +interface BrowserMdxSeparatorOptions { + overlap?: number; + runtimeBatchSize?: number; + timing?: LocalSeparatorTimingCollector; + onProgress?: (progress: LocalSeparatorProgress) => void; + onProviderChange?: (provider: string) => void; +} + +interface SpectrogramPayload { + data: Float32Array; + dims: number[]; + frames?: number; +} + +function packBatchPayloads(payloads: SpectrogramPayload[]): { data: Float32Array; dims: number[]; itemSize: number } { + const frames = payloads[0].dims[3]; + const dimF = payloads[0].dims[2]; + const batch = payloads.length; + const itemSize = 4 * dimF * frames; + const data = new Float32Array(batch * itemSize); + + payloads.forEach((payload, index) => { + data.set(payload.data, index * itemSize); + }); + + return { + data, + dims: [batch, 4, dimF, frames], + itemSize, + }; +} + +function unpackBatchOutput(outputData: Float32Array, batchInfo: { dims: number[]; itemSize: number }): SpectrogramPayload[] { + const results: SpectrogramPayload[] = []; + for (let index = 0; index < batchInfo.dims[0]; index += 1) { + const start = index * batchInfo.itemSize; + const end = start + batchInfo.itemSize; + results.push({ + data: outputData.slice(start, end), + dims: [1, 4, batchInfo.dims[2], batchInfo.dims[3]], + }); + } + return results; +} + +class BrowserMdxSeparator { + private readonly session: ort.InferenceSession; + private readonly runtimeProvider: LocalRuntimeProvider; + private readonly defaults: LocalSeparatorModelConfig['defaults']; + private readonly metadata: LocalSeparatorModelConfig['metadata']; + public onProgress: (progress: LocalSeparatorProgress) => void; + private overlap: number; + private runtimeBatchSize: number; + private readonly enableDenoise: boolean; + private readonly compensate: number; + private readonly primaryStem: string; + private readonly secondaryStem: string; + private readonly nFft: number; + private readonly hopLength: number; + private readonly chunkSize: number; + private readonly trim: number; + private readonly windowCache = createWindowCache(); + private readonly timing: LocalSeparatorTimingCollector; + private dsp: LocalSeparatorCpuDsp | LocalSeparatorGpuDsp; + private dspMode: 'cpu' | 'gpu-hybrid'; + + public static async create( + session: ort.InferenceSession, + runtimeProvider: LocalRuntimeProvider, + config: LocalSeparatorModelConfig, + options: BrowserMdxSeparatorOptions = {}, + ): Promise { + const timing = options.timing ?? new LocalSeparatorTimingCollector('mdx-separation'); + let dsp: LocalSeparatorCpuDsp | LocalSeparatorGpuDsp | null = null; + let dspMode: 'cpu' | 'gpu-hybrid' = 'cpu'; + + if (runtimeProvider === 'webgpu') { + try { + dsp = await timing.measureAsync('dspInit', () => LocalSeparatorGpuDsp.create(config)); + dspMode = 'gpu-hybrid'; + localSeparatorLog('GPU DSP initialized successfully.'); + } catch (error) { + console.warn('[localSeparator] GPU DSP initialization failed, using CPU DSP.', error); + options.onProviderChange?.('cpu/wasm fallback'); + localSeparatorLog('GPU DSP initialization failed. Falling back to CPU DSP.', error); + } + } + + if (!dsp) { + dsp = new LocalSeparatorCpuDsp(config); + localSeparatorLog('Using CPU DSP.'); + } + + return new BrowserMdxSeparator(session, runtimeProvider, config, { + ...options, + dsp, + dspMode, + timing, + }); + } + + private constructor( + session: ort.InferenceSession, + runtimeProvider: LocalRuntimeProvider, + config: LocalSeparatorModelConfig, + options: BrowserMdxSeparatorOptions & { + dsp: LocalSeparatorCpuDsp | LocalSeparatorGpuDsp; + dspMode: 'cpu' | 'gpu-hybrid'; + timing: LocalSeparatorTimingCollector; + }, + ) { + this.session = session; + this.runtimeProvider = runtimeProvider; + this.defaults = config.defaults; + this.metadata = config.metadata; + this.onProgress = options.onProgress ?? (() => {}); + this.overlap = options.overlap ?? this.defaults.overlap; + this.runtimeBatchSize = Math.max(1, options.runtimeBatchSize ?? 2); + this.enableDenoise = this.defaults.enableDenoise; + this.compensate = this.metadata.compensate; + this.primaryStem = this.metadata.primary_stem ?? 'Vocals'; + this.secondaryStem = this.primaryStem === 'Instrumental' ? 'Vocals' : 'Instrumental'; + this.nFft = this.metadata.mdx_n_fft_scale_set; + this.hopLength = this.defaults.hopLength; + this.trim = Math.floor(this.nFft / 2); + this.chunkSize = this.hopLength * (this.defaults.segmentSize - 1); + this.dsp = options.dsp; + this.dspMode = options.dspMode; + this.timing = options.timing; + } + + public dispose(): void { + this.dsp.dispose(); + } + + public getDebugSummary(extra: Record = {}): Record { + return this.timing.getSummary({ + runtimeProvider: this.runtimeProvider, + dspMode: this.dspMode, + ...extra, + }); + } + + public async separate(channels: StereoChannels): Promise<{ + stems: Record; + primaryStem: string; + secondaryStem: string; + }> { + this.onProgress({ stage: 'main', passLabel: 'Main pass', percent: 0, processedChunks: 0, totalChunks: 0 }); + + const { channels: normalizedChannels, originalPeak } = this.timing.measureSync( + 'normalize', + () => normalizeChannels(channels, this.defaults.normalizationThreshold, this.defaults.amplificationThreshold), + ); + const primarySource = await this.demix(normalizedChannels, false); + this.onProgress({ + stage: 'main-complete', + passLabel: 'Main pass', + percent: this.defaults.invertUsingSpec ? 50 : 100, + processedChunks: 0, + totalChunks: 0, + }); + + const primaryScaled = this.timing.measureSync('scalePrimary', () => scaleChannels(primarySource, originalPeak)); + + let secondaryChannels: StereoChannels; + if (this.defaults.invertUsingSpec) { + const rawMix = await this.demix(normalizedChannels, true); + secondaryChannels = this.timing.measureSync('secondaryFromMix', () => { + const secondaryLeft = new Float32Array(rawMix[0].length); + const secondaryRight = new Float32Array(rawMix[1].length); + + for (let i = 0; i < secondaryLeft.length; i += 1) { + secondaryLeft[i] = rawMix[0][i] - (primaryScaled[0][i] * this.compensate); + secondaryRight[i] = rawMix[1][i] - (primaryScaled[1][i] * this.compensate); + } + + return [secondaryLeft, secondaryRight]; + }); + } else { + secondaryChannels = this.timing.measureSync('secondarySubtract', () => { + const secondaryLeft = new Float32Array(primaryScaled[0].length); + const secondaryRight = new Float32Array(primaryScaled[1].length); + for (let i = 0; i < secondaryLeft.length; i += 1) { + secondaryLeft[i] = normalizedChannels[0][i] - (primaryScaled[0][i] * this.compensate); + secondaryRight[i] = normalizedChannels[1][i] - (primaryScaled[1][i] * this.compensate); + } + return [secondaryLeft, secondaryRight]; + }); + } + + const primaryNormalized = this.timing.measureSync( + 'normalizePrimaryOutput', + () => normalizeChannels(primaryScaled, this.defaults.normalizationThreshold, this.defaults.amplificationThreshold).channels, + ); + const secondaryNormalized = this.timing.measureSync( + 'normalizeSecondaryOutput', + () => normalizeChannels(secondaryChannels, this.defaults.normalizationThreshold, this.defaults.amplificationThreshold).channels, + ); + + return { + stems: { + [this.primaryStem]: primaryNormalized, + [this.secondaryStem]: secondaryNormalized, + }, + primaryStem: this.primaryStem, + secondaryStem: this.secondaryStem, + }; + } + + private async demix(channels: StereoChannels, isMatchMix: boolean): Promise { + const overlap = isMatchMix ? this.defaults.matchMixOverlap : this.overlap; + const genSize = this.chunkSize - (2 * this.trim); + const pad = genSize + this.trim - (channels[0].length % genSize); + const mixture: StereoChannels = [ + concatFloat32([new Float32Array(this.trim), channels[0], new Float32Array(pad)]), + concatFloat32([new Float32Array(this.trim), channels[1], new Float32Array(pad)]), + ]; + + const step = Math.max(1, Math.trunc((1 - overlap) * this.chunkSize)); + const result: StereoChannels = [new Float32Array(mixture[0].length), new Float32Array(mixture[1].length)]; + const divider: StereoChannels = [new Float32Array(mixture[0].length), new Float32Array(mixture[1].length)]; + const totalChunks = Math.ceil(mixture[0].length / step); + let processedChunks = 0; + + const windows: Array<{ + start: number; + actualSize: number; + leftChunk: Float32Array; + rightChunk: Float32Array; + window: Float32Array | null; + }> = []; + + for (let start = 0; start < mixture[0].length; start += step) { + const end = Math.min(start + this.chunkSize, mixture[0].length); + const actualSize = end - start; + const leftChunk = new Float32Array(this.chunkSize); + const rightChunk = new Float32Array(this.chunkSize); + leftChunk.set(mixture[0].subarray(start, end)); + rightChunk.set(mixture[1].subarray(start, end)); + windows.push({ + start, + actualSize, + leftChunk, + rightChunk, + window: overlap !== 0 ? getHanning(actualSize, this.windowCache) : null, + }); + } + + for (let batchStart = 0; batchStart < windows.length; batchStart += this.runtimeBatchSize) { + const batch = windows.slice(batchStart, batchStart + this.runtimeBatchSize); + const tarWavesBatch = await this.processBatch(batch, isMatchMix); + + batch.forEach((chunk, index) => { + const tarWaves = tarWavesBatch[index]; + for (let i = 0; i < chunk.actualSize; i += 1) { + const weight = chunk.window ? chunk.window[i] : 1; + result[0][chunk.start + i] += tarWaves[0][i] * weight; + result[1][chunk.start + i] += tarWaves[1][i] * weight; + divider[0][chunk.start + i] += weight; + divider[1][chunk.start + i] += weight; + } + + processedChunks += 1; + const passFraction = totalChunks > 0 ? processedChunks / totalChunks : 1; + const overallPercent = isMatchMix ? 50 + (passFraction * 50) : passFraction * (this.defaults.invertUsingSpec ? 50 : 100); + this.onProgress({ + stage: isMatchMix ? 'match-mix' : 'main', + passLabel: isMatchMix ? 'Match-mix pass' : 'Main pass', + percent: overallPercent, + processedChunks, + totalChunks, + }); + }); + } + + const left = new Float32Array(channels[0].length); + const right = new Float32Array(channels[1].length); + const endTrim = result[0].length - this.trim; + for (let i = this.trim; i < endTrim; i += 1) { + const outIndex = i - this.trim; + if (outIndex >= left.length) break; + left[outIndex] = divider[0][i] > 1e-8 ? result[0][i] / divider[0][i] : 0; + right[outIndex] = divider[1][i] > 1e-8 ? result[1][i] / divider[1][i] : 0; + } + + return [left, right]; + } + + private async processBatch( + batch: Array<{ leftChunk: Float32Array; rightChunk: Float32Array }>, + isMatchMix: boolean, + ): Promise { + const spectra = await this.timing.measureAsync( + isMatchMix ? 'matchMixDspForward' : 'dspForward', + () => Promise.all(batch.map(chunk => this.dsp.forwardStereo(chunk.leftChunk, chunk.rightChunk))), + ); + + if (isMatchMix) { + return this.timing.measureAsync( + 'matchMixDspInverse', + () => Promise.all(spectra.map(payload => this.dsp.inverseStereo(payload))), + ); + } + + let predictedPayloads; + if (this.enableDenoise) { + const positiveOutput = await this.executeModelBatch(spectra); + const negativePayloads = spectra.map(payload => ({ + data: negateArray(payload.data), + dims: payload.dims, + })); + const negativeOutput = await this.executeModelBatch(negativePayloads); + predictedPayloads = positiveOutput.map((payload, index) => { + const data = new Float32Array(payload.data.length); + for (let i = 0; i < data.length; i += 1) { + data[i] = (negativeOutput[index].data[i] * -0.5) + (payload.data[i] * 0.5); + } + return { data, dims: payload.dims }; + }); + } else { + predictedPayloads = await this.executeModelBatch(spectra); + } + + return this.timing.measureAsync( + 'dspInverse', + () => Promise.all(predictedPayloads.map(payload => this.dsp.inverseStereo(payload))), + ); + } + + private async executeModelBatch(payloads: SpectrogramPayload[]): Promise { + const packed = packBatchPayloads(payloads); + const tensor = new ort.Tensor('float32', packed.data, packed.dims); + const feeds = { [this.session.inputNames[0]]: tensor }; + try { + const outputs = await this.timing.measureAsync('inference', () => this.session.run(feeds)); + const firstOutputName = this.session.outputNames[0]; + return unpackBatchOutput(outputs[firstOutputName].data as Float32Array, packed); + } catch (error) { + if (payloads.length > 1 && this.runtimeBatchSize > 1) { + console.warn('[localSeparator] Batched inference failed, falling back to batch size 1.', error); + localSeparatorLog('Batched inference failed. Falling back to batch size 1.', error); + this.runtimeBatchSize = 1; + const singleResults: SpectrogramPayload[] = []; + for (const payload of payloads) { + const singlePacked = packBatchPayloads([payload]); + const singleTensor = new ort.Tensor('float32', singlePacked.data, singlePacked.dims); + const singleFeeds = { [this.session.inputNames[0]]: singleTensor }; + const outputs = await this.timing.measureAsync('inferenceFallback', () => this.session.run(singleFeeds)); + const firstOutputName = this.session.outputNames[0]; + singleResults.push(...unpackBatchOutput(outputs[firstOutputName].data as Float32Array, singlePacked)); + } + return singleResults; + } + throw error; + } + } +} + +export async function decodeAudioToStereo(arrayBuffer: ArrayBuffer): Promise { + const audioContext = new AudioContext({ sampleRate: SAMPLE_RATE }); + const decoded = await audioContext.decodeAudioData(arrayBuffer.slice(0)); + let buffer = decoded; + + if (decoded.sampleRate !== SAMPLE_RATE) { + const offline = new OfflineAudioContext({ + numberOfChannels: Math.max(2, decoded.numberOfChannels), + length: Math.ceil(decoded.duration * SAMPLE_RATE), + sampleRate: SAMPLE_RATE, + }); + const source = offline.createBufferSource(); + source.buffer = decoded; + source.connect(offline.destination); + source.start(); + buffer = await offline.startRendering(); + } + + await audioContext.close(); + + if (buffer.numberOfChannels === 1) { + const mono = buffer.getChannelData(0); + return [new Float32Array(mono), new Float32Array(mono)]; + } + + return [new Float32Array(buffer.getChannelData(0)), new Float32Array(buffer.getChannelData(1))]; +} + +export function channelsToWavBlob(channels: StereoChannels, sampleRate: number = SAMPLE_RATE): Blob { + const length = channels[0].length; + const interleaved = new Int16Array(length * 2); + for (let i = 0; i < length; i += 1) { + interleaved[i * 2] = toInt16(channels[0][i]); + interleaved[(i * 2) + 1] = toInt16(channels[1][i]); + } + + const buffer = new ArrayBuffer(44 + (interleaved.length * 2)); + const view = new DataView(buffer); + writeAscii(view, 0, 'RIFF'); + view.setUint32(4, 36 + (interleaved.length * 2), true); + writeAscii(view, 8, 'WAVE'); + writeAscii(view, 12, 'fmt '); + view.setUint32(16, 16, true); + view.setUint16(20, 1, true); + view.setUint16(22, 2, true); + view.setUint32(24, sampleRate, true); + view.setUint32(28, sampleRate * 4, true); + view.setUint16(32, 4, true); + view.setUint16(34, 16, true); + writeAscii(view, 36, 'data'); + view.setUint32(40, interleaved.length * 2, true); + + let offset = 44; + for (let i = 0; i < interleaved.length; i += 1) { + view.setInt16(offset, interleaved[i], true); + offset += 2; + } + + return new Blob([buffer], { type: 'audio/wav' }); +} + +function writeAscii(view: DataView, offset: number, text: string): void { + for (let i = 0; i < text.length; i += 1) { + view.setUint8(offset + i, text.charCodeAt(i)); + } +} + +function toInt16(value: number): number { + const clamped = Math.max(-1, Math.min(1, value)); + return clamped < 0 ? Math.round(clamped * 0x8000) : Math.round(clamped * 0x7fff); +} + +function concatChannelPairs(chunks: StereoChannels[]): StereoChannels { + return [concatFloat32(chunks.map(chunk => chunk[0])), concatFloat32(chunks.map(chunk => chunk[1]))]; +} + +export async function runLocalSeparator(options: { + session: ort.InferenceSession; + runtimeProvider: LocalRuntimeProvider; + modelConfig: LocalSeparatorModelConfig; + audioBuffer: ArrayBuffer; + chunkDurationSeconds: number | null; + overlap: number; + onProgress: (progress: LocalSeparatorProgress) => void; + onProviderChange?: (provider: string) => void; +}): Promise<{ + stems: Array<{ name: string; blob: Blob }>; + providerLabel: string; + debugSummary: Record; +}> { + const timing = new LocalSeparatorTimingCollector('local-separation'); + const decoded = await timing.measureAsync('decode', () => decodeAudioToStereo(options.audioBuffer)); + localSeparatorLog(`Running browser MDX separation on ${options.runtimeProvider === 'webgpu' ? 'GPU/WebGPU' : 'CPU/wasm'}...`); + + const separator = await BrowserMdxSeparator.create( + options.session, + options.runtimeProvider, + options.modelConfig, + { + overlap: options.overlap, + runtimeBatchSize: options.modelConfig.defaults.batchSize > 1 ? options.modelConfig.defaults.batchSize : 2, + timing, + onProgress: options.onProgress, + onProviderChange: options.onProviderChange, + }, + ); + + try { + const outputs = await separateWithOptionalChunking( + separator, + decoded, + timing, + options.chunkDurationSeconds, + options.modelConfig, + options.onProgress, + ); + const primaryBlob = timing.measureSync('wavEncodePrimary', () => channelsToWavBlob(outputs.stems[outputs.primaryStem])); + const secondaryBlob = timing.measureSync('wavEncodeSecondary', () => channelsToWavBlob(outputs.stems[outputs.secondaryStem])); + + return { + stems: [ + { name: outputs.primaryStem, blob: primaryBlob }, + { name: outputs.secondaryStem, blob: secondaryBlob }, + ], + providerLabel: options.runtimeProvider === 'webgpu' ? 'GPU/WebGPU' : 'CPU/wasm', + debugSummary: separator.getDebugSummary({ model: options.modelConfig.filename }), + }; + } finally { + localSeparatorLog('Separation timing summary', separator.getDebugSummary({ model: options.modelConfig.filename })); + separator.dispose(); + } +} + +async function separateWithOptionalChunking( + separator: BrowserMdxSeparator, + decoded: StereoChannels, + timing: LocalSeparatorTimingCollector, + chunkDurationSeconds: number | null, + modelConfig: LocalSeparatorModelConfig, + onProgress: (progress: LocalSeparatorProgress) => void, +): Promise<{ + stems: Record; + primaryStem: string; + secondaryStem: string; +}> { + if (!chunkDurationSeconds) { + return separator.separate(decoded); + } + + const chunkSamples = Math.max(1, Math.floor(chunkDurationSeconds * SAMPLE_RATE)); + if (decoded[0].length <= chunkSamples) { + return separator.separate(decoded); + } + + const totalChunks = Math.ceil(decoded[0].length / chunkSamples); + const primaryStem = modelConfig.metadata.primary_stem ?? 'Vocals'; + const secondaryStem = primaryStem === 'Instrumental' ? 'Vocals' : 'Instrumental'; + const primaryStemChunks: StereoChannels[] = []; + const secondaryStemChunks: StereoChannels[] = []; + const baseOnProgress = separator.onProgress; + + for (let index = 0; index < totalChunks; index += 1) { + const start = index * chunkSamples; + const end = Math.min(start + chunkSamples, decoded[0].length); + const chunk = sliceChannels(decoded, start, end); + + separator.onProgress = progress => { + const chunkFraction = progress.percent / 100; + const overallPercent = ((index + chunkFraction) / totalChunks) * 100; + baseOnProgress({ + ...progress, + percent: overallPercent, + passLabel: `Audio chunk ${index + 1}/${totalChunks}: ${progress.passLabel}`, + }); + }; + + onProgress({ + stage: 'chunk-prep', + passLabel: `Audio chunk ${index + 1}/${totalChunks}: preparing ${Math.round((end - start) / SAMPLE_RATE)}s chunk...`, + percent: (index / totalChunks) * 100, + processedChunks: index, + totalChunks, + }); + + const result = await timing.measureAsync('chunkedSeparate', () => separator.separate(chunk)); + primaryStemChunks.push(result.stems[primaryStem]); + secondaryStemChunks.push(result.stems[secondaryStem]); + } + + separator.onProgress = baseOnProgress; + + return { + stems: { + [primaryStem]: concatChannelPairs(primaryStemChunks), + [secondaryStem]: concatChannelPairs(secondaryStemChunks), + }, + primaryStem, + secondaryStem, + }; +} diff --git a/src/util/localSeparatorRuntime.ts b/src/util/localSeparatorRuntime.ts new file mode 100644 index 0000000..b44b17c --- /dev/null +++ b/src/util/localSeparatorRuntime.ts @@ -0,0 +1,95 @@ +import * as ort from 'onnxruntime-web/webgpu'; +import ortWasmJsepMjsUrl from 'onnxruntime-web/ort-wasm-simd-threaded.jsep.mjs?url'; +import ortWasmJsepUrl from 'onnxruntime-web/ort-wasm-simd-threaded.jsep.wasm?url'; +import type { LocalRuntimeState, LocalRuntimeSupport, LocalSeparatorModelConfig } from './localSeparatorTypes'; + +function localSeparatorLog(message: string, payload?: unknown): void { + if (payload === undefined) { + console.log(`[localSeparator] ${message}`); + return; + } + console.log(`[localSeparator] ${message}`, payload); +} + +export function detectLocalRuntimeSupport(): LocalRuntimeSupport { + const support = { + webgpuExposed: typeof navigator !== 'undefined' && 'gpu' in navigator, + }; + if (support.webgpuExposed) { + localSeparatorLog('WebGPU API is exposed by this browser.'); + } else { + localSeparatorLog('WebGPU API is not exposed by this browser. CPU/wasm will be used.'); + } + return support; +} + +export class LocalOrtRuntimeManager { + private runtime: LocalRuntimeState | null = null; + private currentModel: string | null = null; + private readonly onProviderChange: (provider: string) => void; + private static wasmPathsConfigured = false; + + constructor({ onProviderChange }: { onProviderChange?: (provider: string) => void } = {}) { + this.onProviderChange = onProviderChange ?? (() => {}); + } + + public reset(): void { + this.runtime = null; + this.currentModel = null; + } + + public async ensureRuntime(modelConfig: LocalSeparatorModelConfig, modelData: Uint8Array): Promise { + if (this.runtime && this.currentModel === modelConfig.filename) { + return this.runtime; + } + + if (!LocalOrtRuntimeManager.wasmPathsConfigured) { + ort.env.wasm.wasmPaths = { + mjs: ortWasmJsepMjsUrl, + wasm: ortWasmJsepUrl, + }; + localSeparatorLog('Configured ONNX Runtime wasm paths.', ort.env.wasm.wasmPaths); + LocalOrtRuntimeManager.wasmPathsConfigured = true; + } + + const providersToTry: Array<'webgpu' | 'wasm'> = []; + if (typeof navigator !== 'undefined' && 'gpu' in navigator) { + providersToTry.push('webgpu'); + localSeparatorLog('navigator.gpu is available, trying WebGPU first.'); + } else { + localSeparatorLog('navigator.gpu is not available. Falling back to CPU/wasm.'); + } + providersToTry.push('wasm'); + + let lastError: unknown = null; + for (const provider of providersToTry) { + try { + if (provider === 'webgpu' && ort.env?.webgpu) { + ort.env.webgpu.powerPreference = 'high-performance'; + localSeparatorLog('Using WebGPU power preference high-performance.'); + } + + const session = await ort.InferenceSession.create(modelData, { + executionProviders: [provider], + graphOptimizationLevel: 'all', + }); + + this.runtime = { provider, session }; + this.currentModel = modelConfig.filename; + this.onProviderChange(provider === 'wasm' ? 'cpu/wasm' : provider); + localSeparatorLog(`Using provider: ${provider}`); + return this.runtime; + } catch (error) { + lastError = error; + if (provider === 'webgpu') { + localSeparatorLog('WebGPU session creation failed. Falling back to CPU/wasm.', error); + this.onProviderChange('cpu/wasm fallback'); + } else { + localSeparatorLog(`Provider failed: ${provider}`, error); + } + } + } + + throw new Error(`Unable to create ONNX Runtime session. ${lastError ? String(lastError) : ''}`.trim()); + } +} diff --git a/src/util/localSeparatorShared.ts b/src/util/localSeparatorShared.ts new file mode 100644 index 0000000..b6549f1 --- /dev/null +++ b/src/util/localSeparatorShared.ts @@ -0,0 +1,275 @@ +import type { StereoChannels } from './localSeparatorTypes'; + +export function index4d(dims: number[], i0: number, i1: number, i2: number, i3: number): number { + return (((i0 * dims[1] + i1) * dims[2] + i2) * dims[3]) + i3; +} + +export function concatFloat32(parts: Float32Array[]): Float32Array { + const length = parts.reduce((sum, part) => sum + part.length, 0); + const output = new Float32Array(length); + let offset = 0; + for (const part of parts) { + output.set(part, offset); + offset += part.length; + } + return output; +} + +export function reflectPad(signal: Float32Array, leftPad: number, rightPad: number): Float32Array { + const result = new Float32Array(leftPad + signal.length + rightPad); + const last = signal.length - 1; + + for (let i = 0; i < leftPad; i += 1) { + result[i] = signal[leftPad - i]; + } + result.set(signal, leftPad); + for (let i = 0; i < rightPad; i += 1) { + result[leftPad + signal.length + i] = signal[last - 1 - i]; + } + + return result; +} + +export function nextPowerOfTwo(value: number): number { + let result = 1; + while (result < value) { + result <<= 1; + } + return result; +} + +export function scaleChannels(channels: StereoChannels, scale: number): StereoChannels { + return channels.map(channel => { + const output = new Float32Array(channel.length); + for (let i = 0; i < channel.length; i += 1) { + output[i] = channel[i] * scale; + } + return output; + }) as StereoChannels; +} + +export function negateArray(data: Float32Array): Float32Array { + const output = new Float32Array(data.length); + for (let i = 0; i < data.length; i += 1) { + output[i] = -data[i]; + } + return output; +} + +export function sliceChannels(channels: StereoChannels, startSample: number, endSample: number): StereoChannels { + return [ + channels[0].slice(startSample, endSample), + channels[1].slice(startSample, endSample), + ]; +} + +export function normalizeChannels( + channels: StereoChannels, + maxPeak: number, + minPeak: number | null, +): { channels: StereoChannels; originalPeak: number } { + let peak = 0; + for (const channel of channels) { + for (let i = 0; i < channel.length; i += 1) { + peak = Math.max(peak, Math.abs(channel[i])); + } + } + + if (peak === 0) { + return { channels, originalPeak: 0 }; + } + + let scale = 1; + if (peak > maxPeak) { + scale = maxPeak / peak; + } else if (minPeak !== null && peak < minPeak && minPeak > 0) { + scale = minPeak / peak; + } + + const normalized = channels.map(channel => { + const output = new Float32Array(channel.length); + for (let i = 0; i < channel.length; i += 1) { + output[i] = channel[i] * scale; + } + return output; + }) as StereoChannels; + + return { channels: normalized, originalPeak: peak }; +} + +export function createWindowCache(): { periodic: Map; symmetric: Map } { + return { + periodic: new Map(), + symmetric: new Map(), + }; +} + +export function getHannPeriodic( + length: number, + cache: { periodic: Map }, +): Float32Array { + const hit = cache.periodic.get(length); + if (hit) return hit; + + const window = new Float32Array(length); + for (let i = 0; i < length; i += 1) { + window[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / length); + } + cache.periodic.set(length, window); + return window; +} + +export function getHanning( + length: number, + cache: { symmetric: Map }, +): Float32Array { + const hit = cache.symmetric.get(length); + if (hit) return hit; + + const window = new Float32Array(length); + if (length === 1) { + window[0] = 1; + } else { + for (let i = 0; i < length; i += 1) { + window[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (length - 1)); + } + } + cache.symmetric.set(length, window); + return window; +} + +export class FFT { + private readonly size: number; + + constructor(size: number) { + this.size = size; + } + + public transform(real: Float64Array, imag: Float64Array): void { + if (real.length !== imag.length || real.length !== this.size) { + throw new Error('FFT input shape mismatch.'); + } + + if ((this.size & (this.size - 1)) === 0) { + this.transformRadix2(real, imag); + } else { + this.transformBluestein(real, imag); + } + } + + public inverse(real: Float64Array, imag: Float64Array): void { + for (let i = 0; i < this.size; i += 1) { + imag[i] = -imag[i]; + } + this.transform(real, imag); + for (let i = 0; i < this.size; i += 1) { + real[i] /= this.size; + imag[i] = -imag[i] / this.size; + } + } + + private transformRadix2(real: Float64Array, imag: Float64Array): void { + const n = this.size; + const levels = Math.trunc(Math.log2(n)); + + for (let i = 0; i < n; i += 1) { + const j = reverseBits(i, levels); + if (j > i) { + [real[i], real[j]] = [real[j], real[i]]; + [imag[i], imag[j]] = [imag[j], imag[i]]; + } + } + + for (let size = 2; size <= n; size <<= 1) { + const halfsize = size >>> 1; + const tableStep = n / size; + for (let i = 0; i < n; i += size) { + for (let j = i, k = 0; j < i + halfsize; j += 1, k += tableStep) { + const angle = (2 * Math.PI * k) / n; + const cos = Math.cos(angle); + const sin = Math.sin(angle); + const tpre = real[j + halfsize] * cos + imag[j + halfsize] * sin; + const tpim = -real[j + halfsize] * sin + imag[j + halfsize] * cos; + real[j + halfsize] = real[j] - tpre; + imag[j + halfsize] = imag[j] - tpim; + real[j] += tpre; + imag[j] += tpim; + } + } + } + } + + private transformBluestein(real: Float64Array, imag: Float64Array): void { + const n = this.size; + const m = nextPowerOfTwo((n * 2) + 1); + const areal = new Float64Array(m); + const aimag = new Float64Array(m); + const breal = new Float64Array(m); + const bimag = new Float64Array(m); + const creal = new Float64Array(m); + const cimag = new Float64Array(m); + + for (let i = 0; i < n; i += 1) { + const angle = (Math.PI * ((i * i) % (n * 2))) / n; + const cos = Math.cos(angle); + const sin = Math.sin(angle); + areal[i] = real[i] * cos + imag[i] * sin; + aimag[i] = -real[i] * sin + imag[i] * cos; + breal[i] = cos; + bimag[i] = sin; + if (i !== 0) { + breal[m - i] = cos; + bimag[m - i] = sin; + } + } + + convolveComplex(areal, aimag, breal, bimag, creal, cimag); + + for (let i = 0; i < n; i += 1) { + const angle = (Math.PI * ((i * i) % (n * 2))) / n; + const cos = Math.cos(angle); + const sin = Math.sin(angle); + real[i] = (creal[i] * cos) + (cimag[i] * sin); + imag[i] = (-creal[i] * sin) + (cimag[i] * cos); + } + } +} + +function reverseBits(x: number, bits: number): number { + let y = 0; + for (let i = 0; i < bits; i += 1) { + y = (y << 1) | (x & 1); + x >>>= 1; + } + return y; +} + +function convolveComplex( + xreal: Float64Array, + ximag: Float64Array, + yreal: Float64Array, + yimag: Float64Array, + outreal: Float64Array, + outimag: Float64Array, +): void { + const n = xreal.length; + const fft = new FFT(n); + const xr = new Float64Array(xreal); + const xi = new Float64Array(ximag); + const yr = new Float64Array(yreal); + const yi = new Float64Array(yimag); + + fft.transform(xr, xi); + fft.transform(yr, yi); + + for (let i = 0; i < n; i += 1) { + const tempReal = (xr[i] * yr[i]) - (xi[i] * yi[i]); + const tempImag = (xi[i] * yr[i]) + (xr[i] * yi[i]); + xr[i] = tempReal; + xi[i] = tempImag; + } + + fft.inverse(xr, xi); + outreal.set(xr); + outimag.set(xi); +} diff --git a/src/util/localSeparatorTiming.ts b/src/util/localSeparatorTiming.ts new file mode 100644 index 0000000..f323a86 --- /dev/null +++ b/src/util/localSeparatorTiming.ts @@ -0,0 +1,52 @@ +export class LocalSeparatorTimingCollector { + private readonly label: string; + private readonly sections = new Map(); + private readonly startedAt = performance.now(); + + constructor(label = 'timing') { + this.label = label; + } + + private track(name: string, durationMs: number): void { + const section = this.sections.get(name) ?? { totalMs: 0, count: 0 }; + section.totalMs += durationMs; + section.count += 1; + this.sections.set(name, section); + } + + public async measureAsync(name: string, fn: () => Promise): Promise { + const started = performance.now(); + try { + return await fn(); + } finally { + this.track(name, performance.now() - started); + } + } + + public measureSync(name: string, fn: () => T): T { + const started = performance.now(); + try { + return fn(); + } finally { + this.track(name, performance.now() - started); + } + } + + public getSummary(extra: Record = {}): Record { + const sections: Record = {}; + for (const [name, section] of this.sections.entries()) { + sections[name] = { + totalMs: Number(section.totalMs.toFixed(2)), + count: section.count, + averageMs: Number((section.totalMs / Math.max(section.count, 1)).toFixed(2)), + }; + } + + return { + label: this.label, + totalMs: Number((performance.now() - this.startedAt).toFixed(2)), + sections, + ...extra, + }; + } +} diff --git a/src/util/localSeparatorTypes.ts b/src/util/localSeparatorTypes.ts new file mode 100644 index 0000000..a39326c --- /dev/null +++ b/src/util/localSeparatorTypes.ts @@ -0,0 +1,49 @@ +export type StereoChannels = [Float32Array, Float32Array]; + +export interface LocalSeparatorModelDefaults { + sampleRate: number; + hopLength: number; + segmentSize: number; + overlap: number; + batchSize: number; + enableDenoise: boolean; + invertUsingSpec: boolean; + normalizationThreshold: number; + amplificationThreshold: number; + matchMixOverlap: number; +} + +export interface LocalSeparatorModelMetadata { + compensate: number; + mdx_dim_f_set: number; + mdx_dim_t_set: number; + mdx_n_fft_scale_set: number; + primary_stem: string; +} + +export interface LocalSeparatorModelConfig { + filename: string; + displayName: string; + status: 'ready'; + defaults: LocalSeparatorModelDefaults; + metadata: LocalSeparatorModelMetadata; +} + +export interface LocalSeparatorProgress { + stage: string; + passLabel: string; + percent: number; + processedChunks: number; + totalChunks: number; +} + +export interface LocalRuntimeSupport { + webgpuExposed: boolean; +} + +export type LocalRuntimeProvider = 'webgpu' | 'wasm'; + +export interface LocalRuntimeState { + provider: LocalRuntimeProvider; + session: import('onnxruntime-web/webgpu').InferenceSession; +} From f57dd7c1dabb6392431fb1d81c0eb870ac55f788 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian <157918347+Xiaohan-Tian@users.noreply.github.com> Date: Mon, 11 May 2026 22:40:13 -0700 Subject: [PATCH 02/19] fix: failed to initialize WebGPU --- src/util/localSeparatorGpuDsp.ts | 18 ++++++++++++++++++ src/util/localSeparatorRunner.ts | 13 ++++++++++--- src/util/localSeparatorRuntime.ts | 6 ++---- 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src/util/localSeparatorGpuDsp.ts b/src/util/localSeparatorGpuDsp.ts index dad4ba0..c3f70dc 100644 --- a/src/util/localSeparatorGpuDsp.ts +++ b/src/util/localSeparatorGpuDsp.ts @@ -2,6 +2,14 @@ import type { LocalSeparatorModelConfig } from './localSeparatorTypes'; import { LocalSeparatorCpuDsp } from './localSeparatorCpuDsp'; import { reflectPad } from './localSeparatorShared'; +function localSeparatorLog(message: string, payload?: unknown): void { + if (payload === undefined) { + console.log(`[localSeparator] ${message}`); + return; + } + console.log(`[localSeparator] ${message}`, payload); +} + type GPUDeviceLike = any; type GPUBufferLike = any; type GPUComputePipelineLike = any; @@ -76,14 +84,22 @@ export class LocalSeparatorGpuDsp { throw new Error('WebGPU is not available for GPU DSP.'); } + localSeparatorLog('Requesting WebGPU adapter for GPU DSP.'); const adapter = await (navigator as { gpu?: { requestAdapter: (options: { powerPreference: string }) => Promise } }).gpu?.requestAdapter({ powerPreference: 'high-performance', }); if (!adapter) { throw new Error('No WebGPU adapter was available for GPU DSP.'); } + localSeparatorLog('WebGPU adapter acquired for GPU DSP.', { + features: typeof adapter.features?.values === 'function' ? Array.from(adapter.features.values()) : undefined, + limits: adapter.limits, + info: typeof adapter.info === 'object' ? adapter.info : undefined, + }); + localSeparatorLog('Requesting WebGPU device for GPU DSP.'); const device = await adapter.requestDevice(); + localSeparatorLog('WebGPU device acquired for GPU DSP.'); return new LocalSeparatorGpuDsp(config, device); } @@ -93,6 +109,7 @@ export class LocalSeparatorGpuDsp { this.nFft = config.metadata.mdx_n_fft_scale_set; this.hopLength = config.defaults.hopLength; this.trim = Math.floor(this.nFft / 2); + localSeparatorLog('Creating GPU DSP compute pipeline.'); this.pipeline = device.createComputePipeline({ layout: 'auto', compute: { @@ -100,6 +117,7 @@ export class LocalSeparatorGpuDsp { entryPoint: 'main', }, }); + localSeparatorLog('GPU DSP compute pipeline created.'); } public async forwardStereo(leftChunk: Float32Array, rightChunk: Float32Array): Promise<{ diff --git a/src/util/localSeparatorRunner.ts b/src/util/localSeparatorRunner.ts index 8419093..f4634b5 100644 --- a/src/util/localSeparatorRunner.ts +++ b/src/util/localSeparatorRunner.ts @@ -111,14 +111,21 @@ class BrowserMdxSeparator { localSeparatorLog('GPU DSP initialized successfully.'); } catch (error) { console.warn('[localSeparator] GPU DSP initialization failed, using CPU DSP.', error); - options.onProviderChange?.('cpu/wasm fallback'); - localSeparatorLog('GPU DSP initialization failed. Falling back to CPU DSP.', error); + options.onProviderChange?.('webgpu + cpu dsp fallback'); + localSeparatorLog( + 'GPU DSP initialization failed. Inference session may still use WebGPU, but DSP will fall back to CPU.', + error, + ); } } if (!dsp) { dsp = new LocalSeparatorCpuDsp(config); - localSeparatorLog('Using CPU DSP.'); + if (runtimeProvider === 'webgpu') { + localSeparatorLog('Using CPU DSP while keeping the WebGPU inference provider.'); + } else { + localSeparatorLog('Using CPU DSP because the active inference provider is CPU/wasm.'); + } } return new BrowserMdxSeparator(session, runtimeProvider, config, { diff --git a/src/util/localSeparatorRuntime.ts b/src/util/localSeparatorRuntime.ts index b44b17c..d01d266 100644 --- a/src/util/localSeparatorRuntime.ts +++ b/src/util/localSeparatorRuntime.ts @@ -1,6 +1,5 @@ import * as ort from 'onnxruntime-web/webgpu'; -import ortWasmJsepMjsUrl from 'onnxruntime-web/ort-wasm-simd-threaded.jsep.mjs?url'; -import ortWasmJsepUrl from 'onnxruntime-web/ort-wasm-simd-threaded.jsep.wasm?url'; +import ortWasmAsyncifyUrl from 'onnxruntime-web/ort-wasm-simd-threaded.asyncify.wasm?url'; import type { LocalRuntimeState, LocalRuntimeSupport, LocalSeparatorModelConfig } from './localSeparatorTypes'; function localSeparatorLog(message: string, payload?: unknown): void { @@ -45,8 +44,7 @@ export class LocalOrtRuntimeManager { if (!LocalOrtRuntimeManager.wasmPathsConfigured) { ort.env.wasm.wasmPaths = { - mjs: ortWasmJsepMjsUrl, - wasm: ortWasmJsepUrl, + wasm: ortWasmAsyncifyUrl, }; localSeparatorLog('Configured ONNX Runtime wasm paths.', ort.env.wasm.wasmPaths); LocalOrtRuntimeManager.wasmPathsConfigured = true; From 64dda5bb0183b67ee1337e32817fec18413baae2 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian <157918347+Xiaohan-Tian@users.noreply.github.com> Date: Mon, 11 May 2026 23:23:11 -0700 Subject: [PATCH 03/19] fix: few minor UI and behavior adjustments --- src/components/KGOnePanel.tsx | 69 +++++++++++++------ src/components/Toolbar.tsx | 24 +++++++ .../settings/sections/GeneralSettings.tsx | 2 +- src/components/track/RegionItem.tsx | 6 +- src/components/track/TrackGridItem.tsx | 1 + src/core/KGProject.ts | 5 +- 6 files changed, 81 insertions(+), 26 deletions(-) diff --git a/src/components/KGOnePanel.tsx b/src/components/KGOnePanel.tsx index 26012dd..a0d48ce 100644 --- a/src/components/KGOnePanel.tsx +++ b/src/components/KGOnePanel.tsx @@ -329,7 +329,7 @@ const ClipTab: React.FC = ({ bpm, keySignature }) => { setGenStatus('polling'); setGenHint('Generating clip...'); - + while (true) { if (signal.aborted) return; @@ -655,7 +655,7 @@ const FullSongTab: React.FC = () => { type ResultItem = { progress: number; stage: string; status: number }; type PollResponse = { data: Array<{ status: number; result: string }>; code: number }; - + while (true) { if (signal.aborted) return; @@ -905,6 +905,19 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => { // eslint-disable-next-line react-hooks/exhaustive-deps }, []); + // Reset separator results whenever the active project changes + useEffect(() => { + stemAudioUrls.forEach(s => URL.revokeObjectURL(s.url)); + setStemAudioUrls([]); + setGenStatus('idle'); + setGenHint(''); + setErrorMsg(''); + setIsImporting(false); + setImportError(''); + originalRegionRef.current = null; + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [projectName]); + useEffect(() => { setModel(availableSeparatorModels[0].value); }, [availableSeparatorModels]); @@ -1116,7 +1129,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => { let files: string[] = []; - + while (true) { if (signal.aborted) return; @@ -1363,7 +1376,8 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
Only Vocal and Instrument (Medium Accuracy) is available while not integrated with K.G.One Music Studio server. Processing in local may take long time depending on your hardware. When fallback to CPU happens, the webpage may - temporarily hang with little or no UI response until processing advances. + temporarily hang with little or no UI response until processing advances.{' '} + Learn more about K.G.One Music Studio server integration.
Provider: {localProviderLabel}
@@ -1740,7 +1754,7 @@ const RemixTab: React.FC = () => { type ResultItem = { progress: number; stage: string; status: number }; type PollResponse = { data: Array<{ status: number; result: string }>; code: number }; - + while (true) { if (signal.aborted) return; @@ -2021,7 +2035,7 @@ const RemixTab: React.FC = () => { const RepaintTab: React.FC = () => { const { selectedRegionIds, projectName, bpm, timeSignature, maxBars, - isLooping, loopingRange, refreshProjectState } = useProjectStore(); + isLooping, loopingRange, refreshProjectState } = useProjectStore(); // Form state (mirrors FullSongTab) const [caption, setCaption] = useState(''); @@ -2258,7 +2272,7 @@ const RepaintTab: React.FC = () => { type ResultItem = { progress: number; stage: string; status: number }; type PollResponse = { data: Array<{ status: number; result: string }>; code: number }; - + while (true) { if (signal.aborted) return; @@ -2573,24 +2587,37 @@ const KGOnePanel: React.FC = ({ isVisible }) => { return (
-

K.G.One Music Generator

+

{mode === 'local-separator' ? 'Music Generator' : 'K.G.One Music Generator'}

{/* Clip tab temporarily disabled, will enable in the future */} - {KGONE_TABS.map(tab => ( - - ))} + {KGONE_TABS.map(tab => { + const isDisabled = disabledTabs.has(tab); + const button = ( + + ); + return isDisabled ? ( + + {button} + + ) : button; + })}
diff --git a/src/components/Toolbar.tsx b/src/components/Toolbar.tsx index b4279be..5381f16 100644 --- a/src/components/Toolbar.tsx +++ b/src/components/Toolbar.tsx @@ -122,6 +122,30 @@ const Toolbar: React.FC = () => { return; } + // Untitled Project is ephemeral — no existing save to preserve, so rename directly + if (savedProjectName === RESERVED_PROJECT_NAME) { + const storage = KGProjectStorage.getInstance(); + const exists = await storage.exists(newName); + if (exists) { + const confirmed = await showConfirm( + `Project "${newName}" already exists. Do you want to overwrite it?` + ); + if (!confirmed) return; + setProjectName(newName); + await saveProject(newName, savedProjectName, setStatus, (finalName) => { + setSavedProjectName(finalName); + if (finalName !== newName) setProjectName(finalName); + }, true /* forceOverwrite */); + return; + } + setProjectName(newName); + await saveProject(newName, savedProjectName, setStatus, (finalName) => { + setSavedProjectName(finalName); + if (finalName !== newName) setProjectName(finalName); + }); + return; + } + // Ask whether the user wants to rename or save as a copy const choice = await showChoice( "Would you like to rename this project, or save it as a new copy?", diff --git a/src/components/settings/sections/GeneralSettings.tsx b/src/components/settings/sections/GeneralSettings.tsx index cbdc30a..8baedfd 100644 --- a/src/components/settings/sections/GeneralSettings.tsx +++ b/src/components/settings/sections/GeneralSettings.tsx @@ -591,7 +591,7 @@ const GeneralSettings: React.FC = () => { disabled={kgoneServerManaged} />
- Base URL of a running K.G.One server. Used for full-song generation, clip generation, and stem separation. + Base URL of a running K.G.One Music Studio server. Used for full-song generation, clip generation, and stem separation.
diff --git a/src/components/track/RegionItem.tsx b/src/components/track/RegionItem.tsx index 9f92bd5..ed5a17a 100644 --- a/src/components/track/RegionItem.tsx +++ b/src/components/track/RegionItem.tsx @@ -45,6 +45,7 @@ interface RegionItemProps { audioBuffer?: AudioBuffer; previewWaveformPeaks?: AudioRecordingPeak[]; isPreview?: boolean; + isAudioRegion?: boolean; } const RegionItem: React.FC = ({ @@ -71,6 +72,7 @@ const RegionItem: React.FC = ({ audioBuffer, previewWaveformPeaks, isPreview = false, + isAudioRegion = false, }) => { // Get selection state and time signature from store const { selectedRegionIds, timeSignature, bpm } = useProjectStore(); @@ -652,7 +654,7 @@ const RegionItem: React.FC = ({ return (
= ({
{name}
-
+
{!isPreview &&
{!audioRegion && (
+ {activeProvider === LOCAL_LLM_PROVIDER_KEY && ( +
+
+

{LOCAL_LLM_DISPLAY_NAME} Local Runtime

+ {!localModelState.runtimeSupport.supported && ( +
+ {localModelState.runtimeSupport.reason} +
+ )} + {!localModelState.isCached && !localModelState.isDownloading && localModelState.runtimeSupport.supported && ( +
+ The local language model has not been downloaded yet. It will be downloaded automatically the next time you send a chat request with this provider. +
+ )} + {(localModelState.isChecking || localModelState.isDownloading || localModelState.progressText) && ( +
+
+
+
+
+ {localModelState.isChecking ? 'Checking local model cache...' : localModelState.progressText} +
+
+ )} + {localModelState.error && ( +
+ {localModelState.error} +
+ )} +
+
+ )} +
{messages.map((message) => ( message.role === 'user' ? ( diff --git a/src/components/settings/Settings.css b/src/components/settings/Settings.css index d916da9..bf85606 100644 --- a/src/components/settings/Settings.css +++ b/src/components/settings/Settings.css @@ -293,6 +293,25 @@ font-size: 12px; } +.settings-progress-block { + margin-top: 10px; +} + +.settings-progress-track { + width: 100%; + height: 10px; + border-radius: 999px; + overflow: hidden; + background-color: #3a3a3a; + border: 1px solid #4a4a4a; +} + +.settings-progress-fill { + height: 100%; + background: linear-gradient(90deg, #5a9fd4 0%, #76c28f 100%); + transition: width 0.2s ease; +} + /* Settings Help Links */ .settings-help-links { display: flex; diff --git a/src/components/settings/sections/GeneralSettings.tsx b/src/components/settings/sections/GeneralSettings.tsx index 8baedfd..8814ad7 100644 --- a/src/components/settings/sections/GeneralSettings.tsx +++ b/src/components/settings/sections/GeneralSettings.tsx @@ -1,8 +1,10 @@ import React, { useState, useEffect, useCallback, useMemo } from 'react'; import { ConfigManager } from '../../../core/config/ConfigManager'; +import { LocalLLMModelManager, type LocalLLMModelState } from '../../../util/localLLMModelManager'; +import { LOCAL_LLM_DISPLAY_NAME, LOCAL_LLM_PROVIDER_KEY } from '../../../util/localLLMConfig'; const GeneralSettings: React.FC = () => { - const [llmProvider, setLlmProvider] = useState('openai'); + const [llmProvider, setLlmProvider] = useState(LOCAL_LLM_PROVIDER_KEY); const [openaiKey, setOpenaiKey] = useState(''); const [openaiModel, setOpenaiModel] = useState(''); const [geminiKey, setGeminiKey] = useState(''); @@ -22,6 +24,7 @@ const GeneralSettings: React.FC = () => { const [kgoneBaseUrl, setKgoneBaseUrl] = useState(''); const [kgoneServerManaged, setKgoneServerManaged] = useState(false); const [soundfontServerManaged, setSoundfontServerManaged] = useState(false); + const [localModelState, setLocalModelState] = useState(LocalLLMModelManager.getState()); const configManager = ConfigManager.instance(); @@ -46,7 +49,7 @@ const GeneralSettings: React.FC = () => { await configManager.initialize(); } - setLlmProvider((configManager.get('general.llm_provider') as string) || 'openai'); + setLlmProvider((configManager.get('general.llm_provider') as string) || LOCAL_LLM_PROVIDER_KEY); setOpenaiKey((configManager.get('general.openai.api_key') as string) || ''); setOpenaiModel((configManager.get('general.openai.model') as string) || ''); setOpenaiFlex((configManager.get('general.openai.flex') as boolean) ?? false); @@ -69,6 +72,8 @@ const GeneralSettings: React.FC = () => { }; loadConfig(); + const unsubscribe = LocalLLMModelManager.subscribe(setLocalModelState); + return unsubscribe; }, [configManager]); // Debounced save function for text inputs @@ -197,6 +202,14 @@ const GeneralSettings: React.FC = () => { debouncedSave('general.kgone.base_url', value); }; + const handleDeleteLocalModel = async () => { + try { + await LocalLLMModelManager.deleteCachedModel(); + } catch (error) { + console.error('Failed to delete local language model cache:', error); + } + }; + // NOTE: Gemini and Claude are not supported yet due to CORS issues. return (
@@ -217,6 +230,7 @@ const GeneralSettings: React.FC = () => { value={llmProvider} onChange={(e) => handleLlmProviderChange(e.target.value)} > + {/* */} @@ -243,6 +257,69 @@ const GeneralSettings: React.FC = () => {
+
+

{LOCAL_LLM_DISPLAY_NAME} Local Runtime

+ + {!localModelState.runtimeSupport.supported && ( +
+ {localModelState.runtimeSupport.reason} +
+ )} + +
+ +
+ {localModelState.isChecking + ? 'Checking local model cache...' + : localModelState.isCached + ? 'Downloaded in browser cache.' + : 'Not downloaded yet.'} +
+
+ + {!localModelState.isCached && !localModelState.isDownloading && localModelState.runtimeSupport.supported && ( +
+ The local model downloads automatically the next time you chat with `Local LLM (Browser)`. +
+ )} + + {(localModelState.isDownloading || localModelState.progressText) && ( +
+
+
+
+
+ {localModelState.progressText} +
+
+ )} + + {localModelState.error && ( +
+ {localModelState.error} +
+ )} + +
+ +
+
+

OpenAI

@@ -600,4 +677,4 @@ const GeneralSettings: React.FC = () => { ); }; -export default GeneralSettings; \ No newline at end of file +export default GeneralSettings; diff --git a/src/constants/coreConstants.ts b/src/constants/coreConstants.ts index f675279..1e03999 100644 --- a/src/constants/coreConstants.ts +++ b/src/constants/coreConstants.ts @@ -104,7 +104,7 @@ export const OPFS_CONSTANTS = { export const CONFIG_UPGRADER_CONSTANTS = { VERSION_KEY: '__config_version', - CURRENT_VERSION: 1, + CURRENT_VERSION: 2, }; export const URL_CONSTANTS = { diff --git a/src/core/config-upgrader/KGConfigUpgrader.ts b/src/core/config-upgrader/KGConfigUpgrader.ts index 63ee19f..22066fb 100644 --- a/src/core/config-upgrader/KGConfigUpgrader.ts +++ b/src/core/config-upgrader/KGConfigUpgrader.ts @@ -1,6 +1,7 @@ import { KGConfigStorage } from '../io/KGConfigStorage'; import { CONFIG_UPGRADER_CONSTANTS } from '../../constants/coreConstants'; import { upgradeConfigToV1 } from './upgradeConfigToV1'; +import { upgradeConfigToV2 } from './upgradeConfigToV2'; /** * KGConfigUpgrader — Orchestrates app-level migrations (e.g., storage backend changes). @@ -33,6 +34,10 @@ export class KGConfigUpgrader { await upgradeConfigToV1(); break; } + case 2: { + await upgradeConfigToV2(); + break; + } default: { throw new Error(`No config upgrader found for version ${nextVersion}`); } diff --git a/src/core/config-upgrader/upgradeConfigToV2.test.ts b/src/core/config-upgrader/upgradeConfigToV2.test.ts new file mode 100644 index 0000000..6f4c8ce --- /dev/null +++ b/src/core/config-upgrader/upgradeConfigToV2.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const configStore = new Map; lastModified: number }>(); + +vi.mock('../io/KGConfigStorage', () => ({ + KGConfigStorage: { + getInstance: () => ({ + getRaw: vi.fn(async (name: string) => configStore.get(name)?.data ?? null), + saveRaw: vi.fn(async (name: string, data: Record) => { + configStore.set(name, { name, data, lastModified: Date.now() }); + }), + }), + }, +})); + +import { upgradeConfigToV2 } from './upgradeConfigToV2'; + +describe('upgradeConfigToV2', () => { + beforeEach(() => { + configStore.clear(); + }); + + it('pins legacy installs without an explicit provider to the old default provider', async () => { + configStore.set('userConfig', { + name: 'userConfig', + data: { + general: { + openai: { api_key: '', model: 'gpt-5.4-mini', flex: false }, + }, + }, + lastModified: Date.now(), + }); + + await upgradeConfigToV2(); + + expect((configStore.get('userConfig')?.data.general as Record).llm_provider).toBe('openai'); + }); + + it('leaves explicit providers unchanged', async () => { + configStore.set('userConfig', { + name: 'userConfig', + data: { + general: { + llm_provider: 'openai_compatible', + }, + }, + lastModified: Date.now(), + }); + + await upgradeConfigToV2(); + + expect((configStore.get('userConfig')?.data.general as Record).llm_provider).toBe('openai_compatible'); + }); +}); diff --git a/src/core/config-upgrader/upgradeConfigToV2.ts b/src/core/config-upgrader/upgradeConfigToV2.ts new file mode 100644 index 0000000..e7a811b --- /dev/null +++ b/src/core/config-upgrader/upgradeConfigToV2.ts @@ -0,0 +1,25 @@ +import { KGConfigStorage } from '../io/KGConfigStorage'; + +const CONFIG_KEY = 'userConfig'; +const LEGACY_DEFAULT_PROVIDER = 'openai'; + +export async function upgradeConfigToV2(): Promise { + const storage = KGConfigStorage.getInstance(); + const rawConfig = await storage.getRaw(CONFIG_KEY); + if (!rawConfig || typeof rawConfig !== 'object') { + return; + } + + const config = rawConfig as Record; + const general = config.general; + if (!general || typeof general !== 'object') { + return; + } + + if ('llm_provider' in (general as Record)) { + return; + } + + (general as Record).llm_provider = LEGACY_DEFAULT_PROVIDER; + await storage.saveRaw(CONFIG_KEY, config); +} diff --git a/src/core/config/ConfigManager.ts b/src/core/config/ConfigManager.ts index b38884f..118e048 100644 --- a/src/core/config/ConfigManager.ts +++ b/src/core/config/ConfigManager.ts @@ -6,7 +6,7 @@ import { KGConfigStorage } from '../io/KGConfigStorage'; interface AppConfig { general: { language: string; - llm_provider: 'openai' | 'gemini' | 'claude' | 'claude_openrouter' | 'openai_compatible'; + llm_provider: 'local_browser' | 'openai' | 'gemini' | 'claude' | 'claude_openrouter' | 'openai_compatible'; persist_api_keys_non_localhost: boolean; openai: { api_key: string; @@ -183,7 +183,7 @@ export class ConfigManager { this.defaultConfig = { general: { language: 'en_us', - llm_provider: 'openai', + llm_provider: 'local_browser', persist_api_keys_non_localhost: false, openai: { api_key: '', diff --git a/src/core/io/LocalSeparatorModelCache.test.ts b/src/core/io/LocalSeparatorModelCache.test.ts index 0642e25..8e96f03 100644 --- a/src/core/io/LocalSeparatorModelCache.test.ts +++ b/src/core/io/LocalSeparatorModelCache.test.ts @@ -9,10 +9,12 @@ class MockWritableFileStream { this.handle = handle; } - async write(content: ArrayBuffer | ArrayBufferView): Promise { - const bytes = content instanceof ArrayBuffer - ? new Uint8Array(content) - : new Uint8Array(content.buffer, content.byteOffset, content.byteLength); + async write(content: ArrayBuffer | ArrayBufferView | string): Promise { + const bytes = typeof content === 'string' + ? new TextEncoder().encode(content) + : content instanceof ArrayBuffer + ? new Uint8Array(content) + : new Uint8Array(content.buffer, content.byteOffset, content.byteLength); this.chunks.push(new Uint8Array(bytes)); } @@ -44,6 +46,8 @@ class MockFileSystemFileHandle { async getFile(): Promise { return { + size: this.content.byteLength, + text: async () => new TextDecoder().decode(this.content), arrayBuffer: async () => this.content.buffer.slice(0), } as unknown as File; } diff --git a/src/util/localLLMConfig.ts b/src/util/localLLMConfig.ts new file mode 100644 index 0000000..577145d --- /dev/null +++ b/src/util/localLLMConfig.ts @@ -0,0 +1,42 @@ +export const LOCAL_LLM_PROVIDER_KEY = 'local_browser'; +export const LOCAL_LLM_MODEL_URL = + 'http://localhost:3000/models/gemma-4-E4B-it-web.task'; +export const LOCAL_LLM_MODEL_FILENAME = 'gemma-4-E4B-it-web.task'; +export const LOCAL_LLM_DISPLAY_NAME = 'Gemma 4 E4B'; +export const LOCAL_LLM_LEGACY_FILENAMES = [ + 'gemma-3n-E4B-it-int4-Web.litertlm', +]; + +export interface LocalLLMRuntimeSupport { + supported: boolean; + webgpuExposed: boolean; + crossOriginIsolated: boolean; + sharedArrayBufferAvailable: boolean; + secureContext: boolean; + reason: string | null; +} + +export function detectLocalLLMRuntimeSupport(): LocalLLMRuntimeSupport { + const secureContext = typeof window !== 'undefined' ? window.isSecureContext : false; + const crossOriginIsolated = typeof window !== 'undefined' ? window.crossOriginIsolated : false; + const sharedArrayBufferAvailable = typeof SharedArrayBuffer !== 'undefined'; + const webgpuExposed = typeof navigator !== 'undefined' && 'gpu' in navigator; + + let reason: string | null = null; + if (!secureContext) { + reason = 'Local browser LLM requires a secure context (HTTPS or localhost).'; + } else if (!crossOriginIsolated || !sharedArrayBufferAvailable) { + reason = 'Local browser LLM requires SharedArrayBuffer support. Ensure COOP/COEP headers are enabled.'; + } else if (!webgpuExposed) { + reason = 'Local browser LLM currently requires a browser with WebGPU support.'; + } + + return { + supported: reason === null, + webgpuExposed, + crossOriginIsolated, + sharedArrayBufferAvailable, + secureContext, + reason, + }; +} diff --git a/src/util/localLLMModelCache.ts b/src/util/localLLMModelCache.ts new file mode 100644 index 0000000..525411d --- /dev/null +++ b/src/util/localLLMModelCache.ts @@ -0,0 +1,92 @@ +import { OpfsModelCache, type ModelDownloadProgress } from './opfsModelCache'; +import { LOCAL_LLM_MODEL_FILENAME } from './localLLMConfig'; + +const cache = new OpfsModelCache({ directoryName: 'models' }); +let writingToCachePromise: Promise | null = null; + +export { type ModelDownloadProgress }; + +export interface CachedModelStreamResult { + reader: ReadableStreamDefaultReader; + totalBytes: number; + fromCache: boolean; + cacheWritePromise: Promise | null; +} + +export class LocalLLMModelCache { + public static async exists(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise { + return cache.exists(filename); + } + + public static async getFile(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise { + return cache.getFile(filename); + } + + public static async getArrayBuffer(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise { + return cache.getArrayBuffer(filename); + } + + public static async delete(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise { + await cache.delete(filename); + } + + public static async loadModelReaderWithCache( + sourceUrl: string, + filename: string = LOCAL_LLM_MODEL_FILENAME, + onProgress?: (progress: ModelDownloadProgress & { fromCache: boolean }) => void, + ): Promise { + if (writingToCachePromise) { + await writingToCachePromise.catch(() => {}); + } + + if (await this.exists(filename)) { + const file = await this.getFile(filename); + onProgress?.({ + receivedBytes: file.size, + totalBytes: file.size, + percent: 100, + fromCache: true, + }); + return { + reader: file.stream().getReader(), + totalBytes: file.size, + fromCache: true, + cacheWritePromise: null, + }; + } + + const response = await fetch(sourceUrl); + if (!response.ok || !response.body) { + throw new Error(`Model download failed (${response.status})`); + } + + const totalBytesHeader = response.headers.get('Content-Length'); + const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : 0; + const [streamForConsumer, streamForCache] = response.body.tee(); + + writingToCachePromise = cache.downloadStream( + streamForCache, + filename, + totalBytes > 0 ? totalBytes : null, + progress => onProgress?.({ ...progress, fromCache: false }), + ); + writingToCachePromise = writingToCachePromise.finally(() => { + writingToCachePromise = null; + }); + + return { + reader: streamForConsumer.getReader(), + totalBytes, + fromCache: false, + cacheWritePromise: writingToCachePromise, + }; + } + + public static async download( + sourceUrl: string, + filename: string = LOCAL_LLM_MODEL_FILENAME, + onProgress?: (progress: ModelDownloadProgress) => void, + ): Promise { + await cache.download(sourceUrl, filename, onProgress); + } +} diff --git a/src/util/localLLMModelManager.ts b/src/util/localLLMModelManager.ts new file mode 100644 index 0000000..9d2405b --- /dev/null +++ b/src/util/localLLMModelManager.ts @@ -0,0 +1,159 @@ +import { + detectLocalLLMRuntimeSupport, + LOCAL_LLM_LEGACY_FILENAMES, + LOCAL_LLM_MODEL_FILENAME, + LOCAL_LLM_MODEL_URL, + type LocalLLMRuntimeSupport, +} from './localLLMConfig'; +import { LocalLLMModelCache } from './localLLMModelCache'; + +export interface LocalLLMModelState { + isCached: boolean; + isChecking: boolean; + isDownloading: boolean; + isDeleting: boolean; + progressPercent: number; + progressText: string; + error: string; + runtimeSupport: LocalLLMRuntimeSupport; +} + +type Listener = (state: LocalLLMModelState) => void; + +export class LocalLLMModelManager { + private static listeners = new Set(); + private static initialized = false; + private static state: LocalLLMModelState = { + isCached: false, + isChecking: false, + isDownloading: false, + isDeleting: false, + progressPercent: 0, + progressText: '', + error: '', + runtimeSupport: detectLocalLLMRuntimeSupport(), + }; + + public static subscribe(listener: Listener): () => void { + this.listeners.add(listener); + listener(this.getState()); + if (!this.initialized) { + this.initialized = true; + void this.refresh(); + } + return () => this.listeners.delete(listener); + } + + public static getState(): LocalLLMModelState { + return { ...this.state, runtimeSupport: { ...this.state.runtimeSupport } }; + } + + public static async refresh(): Promise { + this.setState({ + isChecking: true, + runtimeSupport: detectLocalLLMRuntimeSupport(), + }); + try { + await this.cleanupLegacyEntries(); + const isCached = await LocalLLMModelCache.exists(); + this.setState({ isCached, error: '' }); + } catch (error) { + this.setState({ error: error instanceof Error ? error.message : String(error) }); + } finally { + this.setState({ isChecking: false }); + } + } + + public static async ensureRuntimeSupported(): Promise { + const runtimeSupport = detectLocalLLMRuntimeSupport(); + this.setState({ runtimeSupport }); + if (!runtimeSupport.supported) { + throw new Error(runtimeSupport.reason ?? 'Local browser LLM is not supported in this browser.'); + } + + await this.cleanupLegacyEntries(); + } + + public static async deleteCachedModel(): Promise { + this.setState({ isDeleting: true, error: '' }); + try { + await LocalLLMModelCache.delete(); + await this.cleanupLegacyEntries(); + this.setState({ + isCached: false, + progressPercent: 0, + progressText: '', + }); + } catch (error) { + this.setState({ error: error instanceof Error ? error.message : String(error) }); + throw error; + } finally { + this.setState({ isDeleting: false }); + } + } + + private static setState(partial: Partial): void { + this.state = { + ...this.state, + ...partial, + }; + for (const listener of this.listeners) { + listener(this.getState()); + } + } + + private static async cleanupLegacyEntries(): Promise { + await Promise.all( + LOCAL_LLM_LEGACY_FILENAMES.map(async legacyFilename => { + try { + await LocalLLMModelCache.delete(legacyFilename); + } catch { + // Ignore best-effort legacy cleanup failures. + } + }), + ); + } + + public static notifyLoadStart(fromCache: boolean): void { + this.setState({ + isDownloading: !fromCache, + progressPercent: fromCache ? 100 : 0, + progressText: fromCache ? 'Loading local language model from browser cache...' : 'Downloading local language model...', + error: '', + }); + } + + public static notifyLoadProgress(receivedBytes: number, totalBytes: number | null, fromCache: boolean): void { + const receivedMb = (receivedBytes / (1024 * 1024)).toFixed(1); + const totalMb = totalBytes ? (totalBytes / (1024 * 1024)).toFixed(1) : null; + this.setState({ + isDownloading: !fromCache, + progressPercent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0, + progressText: fromCache + ? 'Loading local language model from browser cache...' + : totalMb + ? `Downloading local language model... ${receivedMb} / ${totalMb} MB` + : `Downloading local language model... ${receivedMb} MB`, + error: '', + }); + } + + public static notifyCacheReady(): void { + this.setState({ + isCached: true, + isDownloading: false, + progressPercent: 100, + progressText: 'Local language model is ready.', + error: '', + }); + } + + public static notifyLoadError(error: unknown): void { + this.setState({ + isDownloading: false, + progressPercent: 0, + progressText: '', + error: error instanceof Error ? error.message : String(error), + }); + } +} diff --git a/src/util/localSeparatorModelCache.ts b/src/util/localSeparatorModelCache.ts index a4990ac..b941cd1 100644 --- a/src/util/localSeparatorModelCache.ts +++ b/src/util/localSeparatorModelCache.ts @@ -1,50 +1,25 @@ import { LOCAL_SEPARATOR_MODEL_FILENAME } from './localSeparatorConfig'; +import { OpfsModelCache, type ModelDownloadProgress } from './opfsModelCache'; -export interface ModelDownloadProgress { - receivedBytes: number; - totalBytes: number | null; - percent: number; -} +const cache = new OpfsModelCache({ directoryName: 'models' }); + +export { type ModelDownloadProgress }; export class LocalSeparatorModelCache { - private static readonly MODELS_DIR = 'models'; - private static readonly TEMP_SUFFIX = '.download'; - public static async exists(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { - try { - const modelsDir = await this.getModelsDir(); - await modelsDir.getFileHandle(filename); - return true; - } catch { - return false; - } + return cache.exists(filename); } public static async getFile(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { - const modelsDir = await this.getModelsDir(); - const fileHandle = await modelsDir.getFileHandle(filename); - return fileHandle.getFile(); + return cache.getFile(filename); } public static async getArrayBuffer(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { - const file = await this.getFile(filename); - return file.arrayBuffer(); + return cache.getArrayBuffer(filename); } public static async delete(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise { - try { - const modelsDir = await this.getModelsDir(); - await modelsDir.removeEntry(filename); - } catch { - // Ignore missing file cleanup. - } - - try { - const modelsDir = await this.getModelsDir(); - await modelsDir.removeEntry(`${filename}${this.TEMP_SUFFIX}`); - } catch { - // Ignore missing temp file cleanup. - } + await cache.delete(filename); } public static async download( @@ -52,75 +27,6 @@ export class LocalSeparatorModelCache { filename: string = LOCAL_SEPARATOR_MODEL_FILENAME, onProgress?: (progress: ModelDownloadProgress) => void, ): Promise { - const response = await fetch(sourceUrl); - if (!response.ok) { - throw new Error(`Model download failed (${response.status})`); - } - - const modelsDir = await this.getModelsDir(); - const tempName = `${filename}${this.TEMP_SUFFIX}`; - await this.delete(filename); - - const tempHandle = await modelsDir.getFileHandle(tempName, { create: true }); - const writable = await tempHandle.createWritable(); - - try { - const totalBytesHeader = response.headers.get('Content-Length'); - const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : null; - - if (!response.body) { - const buffer = await response.arrayBuffer(); - await writable.write(buffer); - onProgress?.({ - receivedBytes: buffer.byteLength, - totalBytes, - percent: 100, - }); - } else { - const reader = response.body.getReader(); - let receivedBytes = 0; - - while (true) { - const { done, value } = await reader.read(); - if (done) break; - if (!value) continue; - await writable.write(value); - receivedBytes += value.byteLength; - onProgress?.({ - receivedBytes, - totalBytes, - percent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0, - }); - } - } - } catch (error) { - await writable.abort(); - await this.delete(filename); - throw error; - } - - await writable.close(); - - const finalHandle = await modelsDir.getFileHandle(filename, { create: true }); - const finalWritable = await finalHandle.createWritable(); - try { - const tempFile = await tempHandle.getFile(); - await finalWritable.write(await tempFile.arrayBuffer()); - await finalWritable.close(); - } catch (error) { - await finalWritable.abort(); - throw error; - } finally { - try { - await modelsDir.removeEntry(tempName); - } catch { - // Ignore temp cleanup errors. - } - } - } - - private static async getModelsDir(): Promise { - const root = await navigator.storage.getDirectory(); - return root.getDirectoryHandle(this.MODELS_DIR, { create: true }); + await cache.download(sourceUrl, filename, onProgress); } } diff --git a/src/util/messageFilter/UserMessageFilter.ts b/src/util/messageFilter/UserMessageFilter.ts index 789a378..b314b18 100644 --- a/src/util/messageFilter/UserMessageFilter.ts +++ b/src/util/messageFilter/UserMessageFilter.ts @@ -2,6 +2,7 @@ import { clearChatHistoryAndUI } from '../chatUtil'; import { useProjectStore } from '../../stores/projectStore'; import { ConfigManager } from '../../core/config/ConfigManager'; import { SystemPrompts } from '../../agent/core/SystemPrompts'; +import { detectLocalLLMRuntimeSupport, LOCAL_LLM_PROVIDER_KEY } from '../localLLMConfig'; export interface UserMessageFilterResult { // Whether to render the user message bubble (div.message-user) @@ -134,7 +135,18 @@ export async function processUserMessage(originalMessage: string): Promise { + try { + const dir = await this.getDir(); + const fileHandle = await dir.getFileHandle(filename); + const sizeHandle = await dir.getFileHandle(this.getSizeFilename(filename)); + const [file, sizeFile] = await Promise.all([fileHandle.getFile(), sizeHandle.getFile()]); + const expectedSize = Number(await sizeFile.text()); + if (!Number.isFinite(expectedSize) || expectedSize <= 0) { + await this.delete(filename); + return false; + } + if (file.size !== expectedSize) { + await this.delete(filename); + return false; + } + return true; + } catch { + return false; + } + } + + public async getFile(filename: string): Promise { + const dir = await this.getDir(); + const handle = await dir.getFileHandle(filename); + const file = await handle.getFile(); + console.log('[opfsModelCache] Opened cached file.', { + filename, + size: file.size, + }); + return file; + } + + public async getArrayBuffer(filename: string): Promise { + const file = await this.getFile(filename); + return file.arrayBuffer(); + } + + public async delete(filename: string): Promise { + const dir = await this.getDir(); + await this.removeIfExists(dir, filename); + await this.removeIfExists(dir, this.getSizeFilename(filename)); + await this.removeIfExists(dir, `${filename}${this.tempSuffix}`); + await this.removeIfExists(dir, `${this.getSizeFilename(filename)}${this.tempSuffix}`); + } + + public async download( + sourceUrl: string, + filename: string, + onProgress?: (progress: ModelDownloadProgress) => void, + ): Promise { + const response = await fetch(sourceUrl); + if (!response.ok) { + throw new Error(`Model download failed (${response.status})`); + } + const totalBytesHeader = response.headers.get('Content-Length'); + const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : null; + if (!response.body) { + throw new Error('Model download response did not include a readable body.'); + } + await this.downloadStream(response.body, filename, totalBytes, onProgress); + } + + public async downloadStream( + stream: ReadableStream, + filename: string, + totalBytes: number | null, + onProgress?: (progress: ModelDownloadProgress) => void, + ): Promise { + const dir = await this.getDir(); + await this.delete(filename); + + const tempFilename = `${filename}${this.tempSuffix}`; + const tempHandle = await dir.getFileHandle(tempFilename, { create: true }); + const tempWritable = await tempHandle.createWritable(); + const reader = stream.getReader(); + let receivedBytes = 0; + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + await tempWritable.write(value); + receivedBytes += value.byteLength; + onProgress?.({ + receivedBytes, + totalBytes, + percent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0, + }); + } + await tempWritable.close(); + + const sizeValue = totalBytes ?? receivedBytes; + if (!Number.isFinite(sizeValue) || sizeValue <= 0) { + throw new Error('Model download did not provide a valid size.'); + } + + console.log(`[opfsModelCache] Finalizing cached model ${filename} from temp file ${tempFilename}.`); + const finalHandle = await dir.getFileHandle(filename, { create: true }); + const finalWritable = await finalHandle.createWritable(); + try { + const tempFile = await tempHandle.getFile(); + console.log('[opfsModelCache] Temp file ready for finalize copy.', { + filename, + tempFilename, + tempSize: tempFile.size, + expectedSize: sizeValue, + }); + await finalWritable.write(tempFile); + await finalWritable.close(); + } catch (error) { + await finalWritable.abort(); + throw error; + } + + const sizeHandle = await dir.getFileHandle(this.getSizeFilename(filename), { create: true }); + const sizeWritable = await sizeHandle.createWritable(); + try { + await sizeWritable.write(String(sizeValue)); + await sizeWritable.close(); + } catch (error) { + await sizeWritable.abort(); + throw error; + } + + onProgress?.({ + receivedBytes: sizeValue, + totalBytes: sizeValue, + percent: 100, + }); + console.log('[opfsModelCache] Cached model finalize completed.', { + filename, + size: sizeValue, + }); + } catch (error) { + try { + await tempWritable.abort(); + } catch { + // Ignore abort cleanup errors. + } + await this.delete(filename); + throw error; + } finally { + await this.removeIfExists(dir, tempFilename); + reader.releaseLock(); + } + } + + private getSizeFilename(filename: string): string { + return `${filename}${this.sizeSuffix}`; + } + + private async getDir(): Promise { + const root = await navigator.storage.getDirectory(); + return root.getDirectoryHandle(this.directoryName, { create: true }); + } + + private async removeIfExists(dir: FileSystemDirectoryHandle, name: string): Promise { + try { + await dir.removeEntry(name); + } catch { + // Ignore missing entry cleanup. + } + } +} diff --git a/vite.config.ts b/vite.config.ts index 2e3d8de..492a9cd 100644 --- a/vite.config.ts +++ b/vite.config.ts @@ -18,7 +18,17 @@ export default defineConfig({ }, server: { host: true, + headers: { + 'Cross-Origin-Opener-Policy': 'same-origin', + 'Cross-Origin-Embedder-Policy': 'require-corp', + }, // IMPORTANT: MAKE SURE TO UPDATE YOUR OS HOSTS FILE TO POINT `testlocal.com` TO YOUR LOCAL IP (e.g. 127.0.0.1). allowedHosts: ['testlocal.com', '.testlocal.com', 'localhost', '127.0.0.1'], }, + preview: { + headers: { + 'Cross-Origin-Opener-Policy': 'same-origin', + 'Cross-Origin-Embedder-Policy': 'require-corp', + }, + }, }); From eb1c7bd2fcfb7475fc20042926f92478f708151e Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:05:13 -0700 Subject: [PATCH 06/19] feat: added performance statistics info when using embedded LLM --- src/agent/core/AgentCore.ts | 5 +++-- src/agent/llm/LocalBrowserLLMProvider.ts | 10 +++++++++- src/agent/llm/StreamingTypes.ts | 8 +++++++- src/components/ChatBox.css | 8 +++++++- src/components/ChatBox.tsx | 1 + src/components/chat/AssistantMessage.tsx | 23 ++++++++++++++++++++++- src/hooks/useStreamProcessor.ts | 9 +++++++-- src/types/projectTypes.ts | 2 ++ 8 files changed, 58 insertions(+), 8 deletions(-) diff --git a/src/agent/core/AgentCore.ts b/src/agent/core/AgentCore.ts index d404c2f..3dd9b7a 100644 --- a/src/agent/core/AgentCore.ts +++ b/src/agent/core/AgentCore.ts @@ -110,6 +110,7 @@ export class AgentCore { let assistantTextContent = ''; const accumulatedToolCalls: ToolCall[] = []; let finishReason = 'stop'; + let performanceInfo: StreamChunk['performanceInfo']; for await (const chunk of this.llmProvider.generateStream(conversationHistory, systemPrompt, tools)) { if (chunk.type === 'text') { @@ -120,6 +121,7 @@ export class AgentCore { accumulatedToolCalls.push(chunk.toolCall); } else if (chunk.type === 'done') { finishReason = chunk.finishReason ?? 'stop'; + performanceInfo = chunk.performanceInfo; } } @@ -163,10 +165,9 @@ export class AgentCore { // LLM finished with text response (stop reason) this.agentState.updateMessage(this.currentAssistantMessageId, assistantTextContent); continueLoop = false; + yield { type: 'done', content: '', finishReason, performanceInfo }; } } - - yield { type: 'done', content: '', finishReason: 'stop' }; } finally { this.currentUserMessageId = null; this.currentAssistantMessageId = null; diff --git a/src/agent/llm/LocalBrowserLLMProvider.ts b/src/agent/llm/LocalBrowserLLMProvider.ts index 9b8b9cb..fb39812 100644 --- a/src/agent/llm/LocalBrowserLLMProvider.ts +++ b/src/agent/llm/LocalBrowserLLMProvider.ts @@ -296,6 +296,14 @@ export class LocalBrowserLLMProvider implements LLMProvider { yield { type: 'tool_call', content: '', toolCall }; } - yield { type: 'done', content: '', finishReason }; + yield { + type: 'done', + content: '', + finishReason, + performanceInfo: { + prefillTps, + generationTps, + }, + }; } } diff --git a/src/agent/llm/StreamingTypes.ts b/src/agent/llm/StreamingTypes.ts index 738ca01..4600a5f 100644 --- a/src/agent/llm/StreamingTypes.ts +++ b/src/agent/llm/StreamingTypes.ts @@ -4,10 +4,16 @@ import type { ToolCall } from '../core/AgentState'; +export interface PerformanceInfo { + prefillTps?: number; + generationTps?: number; +} + export interface StreamChunk { type: 'text' | 'tool_call' | 'tool_result' | 'done'; content: string; toolCall?: ToolCall; toolResult?: { name: string; success: boolean; result: string }; - finishReason?: string; // 'stop' | 'tool_calls' — present on 'done' chunks + performanceInfo?: PerformanceInfo; + finishReason?: string; } diff --git a/src/components/ChatBox.css b/src/components/ChatBox.css index 4f61eec..6f63369 100644 --- a/src/components/ChatBox.css +++ b/src/components/ChatBox.css @@ -226,6 +226,12 @@ font-weight: bold; } +.message-performance-info { + margin-top: 8px; + font-size: 10px; + color: #909090; +} + /* Abort link styling */ .abort-link { background: none !important; @@ -310,4 +316,4 @@ 100% { background-position: -200% 0%; } -} \ No newline at end of file +} diff --git a/src/components/ChatBox.tsx b/src/components/ChatBox.tsx index 2b9940a..4a1cdfd 100644 --- a/src/components/ChatBox.tsx +++ b/src/components/ChatBox.tsx @@ -434,6 +434,7 @@ const ChatBox: React.FC = ({ isVisible }) => { key={message.id} content={message.content} isStreaming={message.isStreaming} + performanceInfo={message.performanceInfo} onAbort={message.isStreaming ? handleAbort : undefined} /> ) diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx index 522ae4a..ad196be 100644 --- a/src/components/chat/AssistantMessage.tsx +++ b/src/components/chat/AssistantMessage.tsx @@ -3,11 +3,13 @@ import ReactMarkdown from 'react-markdown'; import remarkGfm from 'remark-gfm'; import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism'; +import type { PerformanceInfo } from '../../agent/llm/StreamingTypes'; interface AssistantMessageProps { content: string; isStreaming?: boolean; onAbort?: () => void; + performanceInfo?: PerformanceInfo; } // Memoized code component to prevent SyntaxHighlighter re-renders @@ -30,7 +32,19 @@ const CodeComponent = memo(({ inline, className, children, ...props }: any) => { ); }); -const AssistantMessage: React.FC = ({ content, isStreaming, onAbort }) => { +const formatTps = (value?: number): string | null => { + if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) { + return null; + } + + return value.toFixed(1); +}; + +const AssistantMessage: React.FC = ({ content, isStreaming, onAbort, performanceInfo }) => { + const prefillTps = formatTps(performanceInfo?.prefillTps); + const generationTps = formatTps(performanceInfo?.generationTps); + const hasPerformanceInfo = Boolean(prefillTps || generationTps); + const renderContent = () => { // Handle special abort link for streaming messages if (isStreaming && onAbort && content.includes('click here to abort')) { @@ -83,6 +97,13 @@ const AssistantMessage: React.FC = ({ content, isStreamin
{renderContent()} + {hasPerformanceInfo && ( +
+ {prefillTps ? `Prefill: ${prefillTps} t/s` : 'Prefill: -'} + {' · '} + {generationTps ? `Generation: ${generationTps} t/s` : 'Generation: -'} +
+ )}
); diff --git a/src/hooks/useStreamProcessor.ts b/src/hooks/useStreamProcessor.ts index 31f6011..3a6c210 100644 --- a/src/hooks/useStreamProcessor.ts +++ b/src/hooks/useStreamProcessor.ts @@ -38,6 +38,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce let assistantResponse = ''; let tokenCount = 0; let hasTextContent = false; + let performanceInfo: ChatMessage['performanceInfo']; console.log(`------------ ${logPrefix} ------------`); console.log(input); @@ -65,7 +66,8 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce ...msg, content: assistantResponse, isStreaming: false, - tokenCount: undefined + tokenCount: undefined, + performanceInfo })); console.log('------------ ASSISTANT ------------'); @@ -98,19 +100,22 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce assistantResponse = ''; tokenCount = 0; hasTextContent = false; + performanceInfo = undefined; // Create a fresh streaming placeholder for the next LLM response const nextMsg = createStreamingMessage(); currentStreamingId = nextMsg.id; onMessageAdd(nextMsg); } else if (chunk.type === 'done') { + performanceInfo = chunk.performanceInfo; // Finalize the streaming message if (hasTextContent) { onMessageUpdate(currentStreamingId, (msg) => ({ ...msg, content: assistantResponse, isStreaming: false, - tokenCount: undefined + tokenCount: undefined, + performanceInfo })); } else { // No text in final response — remove empty placeholder diff --git a/src/types/projectTypes.ts b/src/types/projectTypes.ts index 3c59d53..b94f5b4 100644 --- a/src/types/projectTypes.ts +++ b/src/types/projectTypes.ts @@ -1,4 +1,5 @@ import { Transform, type TransformFnParams } from 'class-transformer'; +import type { PerformanceInfo } from '../agent/llm/StreamingTypes'; export interface TimeSignature { numerator: number; @@ -11,6 +12,7 @@ export interface ChatMessage { content: string; isStreaming?: boolean; tokenCount?: number; + performanceInfo?: PerformanceInfo; } /** From c7a1630862a63333e261dde824a5aa9331a17f9c Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:12:36 -0700 Subject: [PATCH 07/19] fix: show cache loading progress --- src/agent/llm/LocalBrowserLLMProvider.ts | 6 +--- src/util/localLLMModelCache.ts | 44 ++++++++++++++++++++---- src/util/localLLMModelManager.ts | 12 ++++--- 3 files changed, 45 insertions(+), 17 deletions(-) diff --git a/src/agent/llm/LocalBrowserLLMProvider.ts b/src/agent/llm/LocalBrowserLLMProvider.ts index fb39812..298e5dd 100644 --- a/src/agent/llm/LocalBrowserLLMProvider.ts +++ b/src/agent/llm/LocalBrowserLLMProvider.ts @@ -66,11 +66,7 @@ export class LocalBrowserLLMProvider implements LLMProvider { LOCAL_LLM_MODEL_URL, LOCAL_LLM_MODEL_FILENAME, progress => { - if (progress.fromCache) { - LocalLLMModelManager.notifyLoadStart(true); - } else { - LocalLLMModelManager.notifyLoadProgress(progress.receivedBytes, progress.totalBytes, false); - } + LocalLLMModelManager.notifyLoadProgress(progress.receivedBytes, progress.totalBytes, progress.fromCache); }, ), ]); diff --git a/src/util/localLLMModelCache.ts b/src/util/localLLMModelCache.ts index 525411d..df71f1f 100644 --- a/src/util/localLLMModelCache.ts +++ b/src/util/localLLMModelCache.ts @@ -13,6 +13,42 @@ export interface CachedModelStreamResult { cacheWritePromise: Promise | null; } +const createProgressReader = ( + file: File, + onProgress?: (progress: ModelDownloadProgress & { fromCache: boolean }) => void, +): ReadableStreamDefaultReader => { + const sourceReader = file.stream().getReader(); + let receivedBytes = 0; + + const monitoredStream = new ReadableStream({ + async pull(controller) { + const { done, value } = await sourceReader.read(); + if (done) { + controller.close(); + return; + } + + if (!value) { + return; + } + + receivedBytes += value.byteLength; + onProgress?.({ + receivedBytes, + totalBytes: file.size, + percent: file.size > 0 ? (receivedBytes / file.size) * 100 : 0, + fromCache: true, + }); + controller.enqueue(value); + }, + async cancel(reason) { + await sourceReader.cancel(reason); + }, + }); + + return monitoredStream.getReader(); +}; + export class LocalLLMModelCache { public static async exists(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise { return cache.exists(filename); @@ -41,14 +77,8 @@ export class LocalLLMModelCache { if (await this.exists(filename)) { const file = await this.getFile(filename); - onProgress?.({ - receivedBytes: file.size, - totalBytes: file.size, - percent: 100, - fromCache: true, - }); return { - reader: file.stream().getReader(), + reader: createProgressReader(file, onProgress), totalBytes: file.size, fromCache: true, cacheWritePromise: null, diff --git a/src/util/localLLMModelManager.ts b/src/util/localLLMModelManager.ts index 9d2405b..d55295f 100644 --- a/src/util/localLLMModelManager.ts +++ b/src/util/localLLMModelManager.ts @@ -116,8 +116,8 @@ export class LocalLLMModelManager { public static notifyLoadStart(fromCache: boolean): void { this.setState({ - isDownloading: !fromCache, - progressPercent: fromCache ? 100 : 0, + isDownloading: true, + progressPercent: 0, progressText: fromCache ? 'Loading local language model from browser cache...' : 'Downloading local language model...', error: '', }); @@ -127,10 +127,12 @@ export class LocalLLMModelManager { const receivedMb = (receivedBytes / (1024 * 1024)).toFixed(1); const totalMb = totalBytes ? (totalBytes / (1024 * 1024)).toFixed(1) : null; this.setState({ - isDownloading: !fromCache, + isDownloading: true, progressPercent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0, progressText: fromCache - ? 'Loading local language model from browser cache...' + ? totalMb + ? `Loading local language model from browser cache... ${receivedMb} / ${totalMb} MB` + : `Loading local language model from browser cache... ${receivedMb} MB` : totalMb ? `Downloading local language model... ${receivedMb} / ${totalMb} MB` : `Downloading local language model... ${receivedMb} MB`, @@ -143,7 +145,7 @@ export class LocalLLMModelManager { isCached: true, isDownloading: false, progressPercent: 100, - progressText: 'Local language model is ready.', + progressText: 'Using the local browser model. No external API requests are being sent.', error: '', }); } From 6460d4ffa003898f9af96767e9c75c3bb0741455 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:33:13 -0700 Subject: [PATCH 08/19] fix: initial global spinner randomly hanging until 10s timeout --- src/App.test.tsx | 47 +++++++++++++++++++++++++++++++++++++++++++---- src/App.tsx | 20 ++++++++++++-------- 2 files changed, 55 insertions(+), 12 deletions(-) diff --git a/src/App.test.tsx b/src/App.test.tsx index 4bd3d52..d25e1ec 100644 --- a/src/App.test.tsx +++ b/src/App.test.tsx @@ -2,6 +2,8 @@ import { act, render, screen } from '@testing-library/react'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; let mockState = { isPreparingPlayback: false }; +let mockActiveLoadCount = 0; +let loadingListener: ((evt: { type: 'start' | 'end'; instrument: string }) => void) | null = null; vi.mock('./stores/projectStore', () => ({ useProjectStore: (selector?: (state: typeof mockState) => unknown) => ( @@ -21,12 +23,21 @@ vi.mock('./components/ChatBox', () => ({ default: () => null })); vi.mock('./components/KGOnePanel', () => ({ default: () => null })); vi.mock('./components/EventListPanel', () => ({ default: () => null })); vi.mock('./components/settings', () => ({ SettingsPanel: () => null })); +vi.mock('./util/dialogUtil', () => ({ + showAlert: vi.fn().mockResolvedValue(undefined), +})); vi.mock('./core/audio-interface/KGToneBuffersPool', () => ({ KGToneBuffersPool: { instance: () => ({ - getActiveLoadCount: () => 0, - addLoadingListener: () => undefined, - removeLoadingListener: () => undefined, + getActiveLoadCount: () => mockActiveLoadCount, + addLoadingListener: (listener: (evt: { type: 'start' | 'end'; instrument: string }) => void) => { + loadingListener = listener; + }, + removeLoadingListener: (listener: (evt: { type: 'start' | 'end'; instrument: string }) => void) => { + if (loadingListener === listener) { + loadingListener = null; + } + }, }), }, })); @@ -47,12 +58,14 @@ vi.mock('./core/KGCore', () => ({ }, })); -import { PlaybackPreparationOverlayContainer } from './App'; +import { GlobalLoadingOverlayContainer, PlaybackPreparationOverlayContainer } from './App'; describe('PlaybackPreparationOverlayContainer', () => { beforeEach(() => { vi.useFakeTimers(); mockState = { isPreparingPlayback: false }; + mockActiveLoadCount = 0; + loadingListener = null; }); afterEach(() => { @@ -112,3 +125,29 @@ describe('PlaybackPreparationOverlayContainer', () => { expect(screen.queryByText('Preparing playback...')).not.toBeInTheDocument(); }); }); + +describe('GlobalLoadingOverlayContainer', () => { + beforeEach(() => { + vi.useFakeTimers(); + mockActiveLoadCount = 0; + loadingListener = null; + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('syncs to the pool active count instead of relying on incremental listener math', () => { + mockActiveLoadCount = 2; + render(); + + expect(screen.getByText('Loading ... (2)')).toBeInTheDocument(); + + mockActiveLoadCount = 0; + act(() => { + loadingListener?.({ type: 'end', instrument: 'woodblock' }); + }); + + expect(screen.queryByText(/Loading \.\.\./)).not.toBeInTheDocument(); + }); +}); diff --git a/src/App.tsx b/src/App.tsx index c31919b..a4e4c9e 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -195,19 +195,20 @@ function App() { export default App; // Local component to subscribe to pool events and manage a counter -const GlobalLoadingOverlayContainer: React.FC = () => { +export const GlobalLoadingOverlayContainer: React.FC = () => { const [loadingCount, setLoadingCount] = useState(() => KGToneBuffersPool.instance().getActiveLoadCount()); const [overdue, setOverdue] = useState(false); const timeoutRef = useRef(null); useEffectReact(() => { const pool = KGToneBuffersPool.instance(); - const listener = (evt: { type: 'start' | 'end'; instrument: string }) => { - setLoadingCount(prev => { - if (evt.type === 'start') return prev + 1; - return Math.max(0, prev - 1); - }); + const syncLoadingCount = () => { + setLoadingCount(pool.getActiveLoadCount()); }; + const listener = (_evt: { type: 'start' | 'end'; instrument: string }) => { + syncLoadingCount(); + }; + syncLoadingCount(); pool.addLoadingListener(listener); return () => { pool.removeLoadingListener(listener); @@ -219,8 +220,11 @@ const GlobalLoadingOverlayContainer: React.FC = () => { // When loading starts, start a 30s timer if not already overdue/timed if (loadingCount > 0 && !overdue && timeoutRef.current === null) { timeoutRef.current = window.setTimeout(async () => { - // Only trigger if still loading - if (loadingCount > 0) { + const activeLoadCount = KGToneBuffersPool.instance().getActiveLoadCount(); + setLoadingCount(activeLoadCount); + + // Only trigger if still loading according to the source of truth + if (activeLoadCount > 0) { setOverdue(true); await showAlert( 'Loading resources is taking longer than expected and may have partially failed. If you notice any playback issues, please refresh the page to retry downloading the audio files.' From dbf0c20542e962c130ec4e9660350b1264dd0ebf Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:33:44 -0700 Subject: [PATCH 09/19] fix: updated Gemma 4 URL from local test url to HF url --- src/util/localLLMConfig.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/util/localLLMConfig.ts b/src/util/localLLMConfig.ts index 577145d..e2f632d 100644 --- a/src/util/localLLMConfig.ts +++ b/src/util/localLLMConfig.ts @@ -1,6 +1,6 @@ export const LOCAL_LLM_PROVIDER_KEY = 'local_browser'; export const LOCAL_LLM_MODEL_URL = - 'http://localhost:3000/models/gemma-4-E4B-it-web.task'; + 'https://huggingface.co/notabilia/gemma-4-E4B-it-litert-lm/resolve/main/gemma-4-E4B-it-web.task'; export const LOCAL_LLM_MODEL_FILENAME = 'gemma-4-E4B-it-web.task'; export const LOCAL_LLM_DISPLAY_NAME = 'Gemma 4 E4B'; export const LOCAL_LLM_LEGACY_FILENAMES = [ From a48966d7615ca141966d6d9af32dec9196952de9 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:44:25 -0700 Subject: [PATCH 10/19] feat: added local LLM context length option --- public/config.json | 3 + src/agent/llm/LocalBrowserLLMProvider.test.ts | 106 ++++++++++++++++ src/agent/llm/LocalBrowserLLMProvider.ts | 24 +++- src/components/ChatBox.tsx | 1 + .../sections/GeneralSettings.test.tsx | 120 ++++++++++++++++++ .../settings/sections/GeneralSettings.tsx | 45 ++++++- src/constants/coreConstants.ts | 2 +- src/core/config-upgrader/KGConfigUpgrader.ts | 5 + .../config-upgrader/upgradeConfigToV3.test.ts | 88 +++++++++++++ src/core/config-upgrader/upgradeConfigToV3.ts | 35 +++++ src/core/config/ConfigManager.ts | 6 + src/util/localLLMConfig.test.ts | 19 +++ src/util/localLLMConfig.ts | 17 +++ 13 files changed, 467 insertions(+), 4 deletions(-) create mode 100644 src/agent/llm/LocalBrowserLLMProvider.test.ts create mode 100644 src/components/settings/sections/GeneralSettings.test.tsx create mode 100644 src/core/config-upgrader/upgradeConfigToV3.test.ts create mode 100644 src/core/config-upgrader/upgradeConfigToV3.ts create mode 100644 src/util/localLLMConfig.test.ts diff --git a/public/config.json b/public/config.json index 4d2f9ea..015918a 100644 --- a/public/config.json +++ b/public/config.json @@ -25,6 +25,9 @@ "api_key": "", "base_url": "", "model": "" + }, + "local_browser": { + "context_length": 32768 }, "soundfont": { "base_url": "https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/" diff --git a/src/agent/llm/LocalBrowserLLMProvider.test.ts b/src/agent/llm/LocalBrowserLLMProvider.test.ts new file mode 100644 index 0000000..1284c9d --- /dev/null +++ b/src/agent/llm/LocalBrowserLLMProvider.test.ts @@ -0,0 +1,106 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +const { + configGetMock, + ensureRuntimeSupportedMock, + notifyLoadProgressMock, + notifyLoadStartMock, + notifyCacheReadyMock, + notifyLoadErrorMock, + loadModelReaderWithCacheMock, +} = vi.hoisted(() => ({ + configGetMock: vi.fn(), + ensureRuntimeSupportedMock: vi.fn(async () => undefined), + notifyLoadProgressMock: vi.fn(), + notifyLoadStartMock: vi.fn(), + notifyCacheReadyMock: vi.fn(), + notifyLoadErrorMock: vi.fn(), + loadModelReaderWithCacheMock: vi.fn(), +})); + +vi.mock('../../core/config/ConfigManager', () => ({ + ConfigManager: { + instance: () => ({ + get: configGetMock, + }), + }, +})); + +vi.mock('../../util/localLLMModelManager', () => ({ + LocalLLMModelManager: { + ensureRuntimeSupported: ensureRuntimeSupportedMock, + notifyLoadProgress: notifyLoadProgressMock, + notifyLoadStart: notifyLoadStartMock, + notifyCacheReady: notifyCacheReadyMock, + notifyLoadError: notifyLoadErrorMock, + }, +})); + +vi.mock('../../util/localLLMModelCache', () => ({ + LocalLLMModelCache: { + loadModelReaderWithCache: loadModelReaderWithCacheMock, + }, +})); + +import { LocalBrowserLLMProvider } from './LocalBrowserLLMProvider'; + +describe('LocalBrowserLLMProvider', () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + beforeEach(() => { + configGetMock.mockReset(); + ensureRuntimeSupportedMock.mockClear(); + notifyLoadProgressMock.mockClear(); + notifyLoadStartMock.mockClear(); + notifyCacheReadyMock.mockClear(); + notifyLoadErrorMock.mockClear(); + loadModelReaderWithCacheMock.mockReset(); + }); + + async function runProviderAndCaptureOptions(configValue: unknown): Promise> { + const createFromOptionsMock = vi.fn(async (_fileset: unknown, options: Record) => ({ + generateResponse: (_prompt: string, callback: (partial: string, done: boolean) => void) => { + callback('hello', true); + }, + sizeInTokens: (text: string) => text.length, + })); + + configGetMock.mockReturnValue(configValue); + loadModelReaderWithCacheMock.mockResolvedValue({ + reader: new Uint8Array([1, 2, 3]), + totalBytes: 3, + fromCache: true, + cacheWritePromise: null, + }); + + vi.spyOn(LocalBrowserLLMProvider.prototype as never, 'getMediaPipeModule' as never).mockResolvedValue({ + FilesetResolver: { + forGenAiTasks: vi.fn(async () => ({})), + }, + LlmInference: { + createFromOptions: createFromOptionsMock, + }, + }); + + const provider = new LocalBrowserLLMProvider(); + const chunks: unknown[] = []; + for await (const chunk of provider.generateStream([])) { + chunks.push(chunk); + } + + expect(chunks.length).toBeGreaterThan(0); + return createFromOptionsMock.mock.calls[0][1] as Record; + } + + it('uses the configured maxTokens value', async () => { + const options = await runProviderAndCaptureOptions(65536); + expect(options.maxTokens).toBe(65536); + }); + + it('falls back to 32768 when config is invalid', async () => { + const options = await runProviderAndCaptureOptions(99999); + expect(options.maxTokens).toBe(32768); + }); +}); diff --git a/src/agent/llm/LocalBrowserLLMProvider.ts b/src/agent/llm/LocalBrowserLLMProvider.ts index 298e5dd..57970f8 100644 --- a/src/agent/llm/LocalBrowserLLMProvider.ts +++ b/src/agent/llm/LocalBrowserLLMProvider.ts @@ -9,9 +9,15 @@ import { parseToolCalls, stripToolProtocol, } from './gemmaToolProtocol'; -import { LOCAL_LLM_MODEL_FILENAME, LOCAL_LLM_MODEL_URL } from '../../util/localLLMConfig'; +import { + LOCAL_LLM_DEFAULT_CONTEXT_LENGTH, + LOCAL_LLM_MODEL_FILENAME, + LOCAL_LLM_MODEL_URL, + normalizeLocalLLMContextLength, +} from '../../util/localLLMConfig'; import { LocalLLMModelCache } from '../../util/localLLMModelCache'; import type { LLMProvider } from './LLMProvider'; +import { ConfigManager } from '../../core/config/ConfigManager'; type MediaPipeGenAI = { FilesetResolver: { @@ -60,6 +66,9 @@ export class LocalBrowserLLMProvider implements LLMProvider { return this.inference; } + const maxTokens = this.getConfiguredContextLength(); + console.log(`[localLLM] Initializing with max context length: ${maxTokens} tokens`); + const [{ FilesetResolver, LlmInference }, modelLoad] = await Promise.all([ this.getMediaPipeModule(), LocalLLMModelCache.loadModelReaderWithCache( @@ -86,7 +95,7 @@ export class LocalBrowserLLMProvider implements LLMProvider { modelAssetBuffer: modelLoad.reader, }, numResponses: 1, - maxTokens: 32768, + maxTokens, topK: 64, temperature: 1.0, }); @@ -112,6 +121,17 @@ export class LocalBrowserLLMProvider implements LLMProvider { return importMediaPipe(); } + private getConfiguredContextLength(): number { + try { + const configManager = ConfigManager.instance(); + return normalizeLocalLLMContextLength( + configManager.get('general.local_browser.context_length'), + ); + } catch { + return LOCAL_LLM_DEFAULT_CONTEXT_LENGTH; + } + } + private applyTemplate(message: { role: 'user' | 'model'; text: string }): string { const template = PROMPT_TEMPLATE[message.role]; return `${template.pre}${message.text}${template.post}`; diff --git a/src/components/ChatBox.tsx b/src/components/ChatBox.tsx index 4a1cdfd..7be2842 100644 --- a/src/components/ChatBox.tsx +++ b/src/components/ChatBox.tsx @@ -200,6 +200,7 @@ const ChatBox: React.FC = ({ isVisible }) => { const unsubscribe = configManager.addChangeListener((changedKeys) => { if ( changedKeys.includes('general.llm_provider') || + changedKeys.includes('general.local_browser.context_length') || changedKeys.some(k => k.startsWith('general.openai.')) || changedKeys.some(k => k.startsWith('general.openai_compatible.')) ) { diff --git a/src/components/settings/sections/GeneralSettings.test.tsx b/src/components/settings/sections/GeneralSettings.test.tsx new file mode 100644 index 0000000..d45197f --- /dev/null +++ b/src/components/settings/sections/GeneralSettings.test.tsx @@ -0,0 +1,120 @@ +import React from 'react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import GeneralSettings from './GeneralSettings'; + +const configState = new Map([ + ['general.llm_provider', 'local_browser'], + ['general.persist_api_keys_non_localhost', false], + ['general.openai.api_key', ''], + ['general.openai.model', 'gpt-5.4-mini'], + ['general.openai.flex', false], + ['general.gemini.api_key', ''], + ['general.gemini.model', 'gemini-2.5-flash'], + ['general.claude.api_key', ''], + ['general.claude.model', 'claude-sonnet-4.6'], + ['general.claude_openrouter.api_key', ''], + ['general.claude_openrouter.base_url', 'https://openrouter.ai/api/v1'], + ['general.claude_openrouter.model', 'anthropic/claude-sonnet-4.6'], + ['general.openai_compatible.api_key', ''], + ['general.openai_compatible.base_url', ''], + ['general.openai_compatible.model', ''], + ['general.local_browser.context_length', 65536], + ['general.soundfont.base_url', 'https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/'], + ['general.kgone.enabled', false], + ['general.kgone.base_url', 'http://127.0.0.1:8000'], +]); + +const configManagerMock = { + getIsInitialized: vi.fn(() => true), + initialize: vi.fn().mockResolvedValue(undefined), + get: vi.fn((key: string) => configState.get(key)), + set: vi.fn(async (key: string, value: unknown) => { + configState.set(key, value); + }), + isKGOneServerManaged: vi.fn(() => false), + isSoundfontServerManaged: vi.fn(() => false), +}; + +vi.mock('../../../core/config/ConfigManager', () => ({ + ConfigManager: { + instance: () => configManagerMock, + }, +})); + +vi.mock('../../../util/localLLMModelManager', () => ({ + LocalLLMModelManager: { + getState: () => ({ + isCached: false, + isChecking: false, + isDownloading: false, + isDeleting: false, + progressPercent: 0, + progressText: '', + error: '', + runtimeSupport: { + supported: true, + webgpuExposed: true, + crossOriginIsolated: true, + sharedArrayBufferAvailable: true, + secureContext: true, + reason: null, + }, + }), + subscribe: (listener: (state: unknown) => void) => { + listener({ + isCached: false, + isChecking: false, + isDownloading: false, + isDeleting: false, + progressPercent: 0, + progressText: '', + error: '', + runtimeSupport: { + supported: true, + webgpuExposed: true, + crossOriginIsolated: true, + sharedArrayBufferAvailable: true, + secureContext: true, + reason: null, + }, + }); + return () => {}; + }, + deleteCachedModel: vi.fn().mockResolvedValue(undefined), + }, +})); + +describe('GeneralSettings', () => { + beforeEach(() => { + configState.set('general.local_browser.context_length', 65536); + configManagerMock.get.mockClear(); + configManagerMock.set.mockClear(); + }); + + it('renders the local context length selector and VRAM hint', async () => { + render(); + + expect(await screen.findByText('Gemma 4 E4B Local Runtime')).toBeTruthy(); + expect(screen.getByLabelText('Context Length')).toBeTruthy(); + expect(screen.getByText(/require more VRAM/i)).toBeTruthy(); + }); + + it('initializes the local context length from config', async () => { + render(); + + const select = await screen.findByLabelText('Context Length'); + expect((select as HTMLSelectElement).value).toBe('65536'); + }); + + it('persists local context length changes', async () => { + render(); + + const select = await screen.findByLabelText('Context Length'); + fireEvent.change(select, { target: { value: '131072' } }); + + await waitFor(() => { + expect(configManagerMock.set).toHaveBeenCalledWith('general.local_browser.context_length', 131072); + }); + }); +}); diff --git a/src/components/settings/sections/GeneralSettings.tsx b/src/components/settings/sections/GeneralSettings.tsx index 8814ad7..b3edf73 100644 --- a/src/components/settings/sections/GeneralSettings.tsx +++ b/src/components/settings/sections/GeneralSettings.tsx @@ -1,7 +1,15 @@ import React, { useState, useEffect, useCallback, useMemo } from 'react'; import { ConfigManager } from '../../../core/config/ConfigManager'; import { LocalLLMModelManager, type LocalLLMModelState } from '../../../util/localLLMModelManager'; -import { LOCAL_LLM_DISPLAY_NAME, LOCAL_LLM_PROVIDER_KEY } from '../../../util/localLLMConfig'; +import { + formatLocalLLMContextLength, + LOCAL_LLM_CONTEXT_LENGTH_OPTIONS, + LOCAL_LLM_DEFAULT_CONTEXT_LENGTH, + LOCAL_LLM_DISPLAY_NAME, + LOCAL_LLM_PROVIDER_KEY, + normalizeLocalLLMContextLength, + type LocalLLMContextLength, +} from '../../../util/localLLMConfig'; const GeneralSettings: React.FC = () => { const [llmProvider, setLlmProvider] = useState(LOCAL_LLM_PROVIDER_KEY); @@ -24,6 +32,7 @@ const GeneralSettings: React.FC = () => { const [kgoneBaseUrl, setKgoneBaseUrl] = useState(''); const [kgoneServerManaged, setKgoneServerManaged] = useState(false); const [soundfontServerManaged, setSoundfontServerManaged] = useState(false); + const [localContextLength, setLocalContextLength] = useState(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH); const [localModelState, setLocalModelState] = useState(LocalLLMModelManager.getState()); const configManager = ConfigManager.instance(); @@ -64,6 +73,7 @@ const GeneralSettings: React.FC = () => { setCompatibleKey((configManager.get('general.openai_compatible.api_key') as string) || ''); setCompatibleBaseUrl((configManager.get('general.openai_compatible.base_url') as string) || ''); setCompatibleModel((configManager.get('general.openai_compatible.model') as string) || ''); + setLocalContextLength(normalizeLocalLLMContextLength(configManager.get('general.local_browser.context_length'))); setSoundfontBaseUrl((configManager.get('general.soundfont.base_url') as string) || ''); setKgoneEnabled((configManager.get('general.kgone.enabled') as boolean) ?? false); setKgoneBaseUrl((configManager.get('general.kgone.base_url') as string) || ''); @@ -210,6 +220,18 @@ const GeneralSettings: React.FC = () => { } }; + const handleLocalContextLengthChange = async (value: string) => { + const parsed = Number(value); + const normalized = normalizeLocalLLMContextLength(parsed); + setLocalContextLength(normalized); + try { + await configManager.set('general.local_browser.context_length', normalized); + console.log('Local browser context length changed to:', normalized); + } catch (error) { + console.error('Failed to save local browser context length:', error); + } + }; + // NOTE: Gemini and Claude are not supported yet due to CORS issues. return (
@@ -279,6 +301,27 @@ const GeneralSettings: React.FC = () => {
+
+ + +
+ Larger context lengths require more VRAM and may also reduce performance as conversations become longer. +
+
+ {!localModelState.isCached && !localModelState.isDownloading && localModelState.runtimeSupport.supported && (
The local model downloads automatically the next time you chat with `Local LLM (Browser)`. diff --git a/src/constants/coreConstants.ts b/src/constants/coreConstants.ts index 1e03999..33f0f40 100644 --- a/src/constants/coreConstants.ts +++ b/src/constants/coreConstants.ts @@ -104,7 +104,7 @@ export const OPFS_CONSTANTS = { export const CONFIG_UPGRADER_CONSTANTS = { VERSION_KEY: '__config_version', - CURRENT_VERSION: 2, + CURRENT_VERSION: 3, }; export const URL_CONSTANTS = { diff --git a/src/core/config-upgrader/KGConfigUpgrader.ts b/src/core/config-upgrader/KGConfigUpgrader.ts index 22066fb..21e08d7 100644 --- a/src/core/config-upgrader/KGConfigUpgrader.ts +++ b/src/core/config-upgrader/KGConfigUpgrader.ts @@ -2,6 +2,7 @@ import { KGConfigStorage } from '../io/KGConfigStorage'; import { CONFIG_UPGRADER_CONSTANTS } from '../../constants/coreConstants'; import { upgradeConfigToV1 } from './upgradeConfigToV1'; import { upgradeConfigToV2 } from './upgradeConfigToV2'; +import { upgradeConfigToV3 } from './upgradeConfigToV3'; /** * KGConfigUpgrader — Orchestrates app-level migrations (e.g., storage backend changes). @@ -38,6 +39,10 @@ export class KGConfigUpgrader { await upgradeConfigToV2(); break; } + case 3: { + await upgradeConfigToV3(); + break; + } default: { throw new Error(`No config upgrader found for version ${nextVersion}`); } diff --git a/src/core/config-upgrader/upgradeConfigToV3.test.ts b/src/core/config-upgrader/upgradeConfigToV3.test.ts new file mode 100644 index 0000000..ea64dfb --- /dev/null +++ b/src/core/config-upgrader/upgradeConfigToV3.test.ts @@ -0,0 +1,88 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const configStore = new Map; lastModified: number }>(); + +vi.mock('../io/KGConfigStorage', () => ({ + KGConfigStorage: { + getInstance: () => ({ + getRaw: vi.fn(async (name: string) => configStore.get(name)?.data ?? null), + saveRaw: vi.fn(async (name: string, data: Record) => { + configStore.set(name, { name, data, lastModified: Date.now() }); + }), + }), + }, +})); + +import { upgradeConfigToV3 } from './upgradeConfigToV3'; + +describe('upgradeConfigToV3', () => { + beforeEach(() => { + configStore.clear(); + }); + + it('adds the default local browser context length when missing', async () => { + configStore.set('userConfig', { + name: 'userConfig', + data: { + general: { + llm_provider: 'local_browser', + }, + }, + lastModified: Date.now(), + }); + + await upgradeConfigToV3(); + + expect( + ((configStore.get('userConfig')?.data.general as Record).local_browser as Record).context_length, + ).toBe(32768); + }); + + it.each([32768, 65536, 131072])('preserves existing context length %s', async (existingValue) => { + configStore.set('userConfig', { + name: 'userConfig', + data: { + general: { + local_browser: { + context_length: existingValue, + }, + }, + }, + lastModified: Date.now(), + }); + + await upgradeConfigToV3(); + + expect( + ((configStore.get('userConfig')?.data.general as Record).local_browser as Record).context_length, + ).toBe(existingValue); + }); + + it('is a no-op when general is missing', async () => { + configStore.set('userConfig', { + name: 'userConfig', + data: {}, + lastModified: Date.now(), + }); + + await upgradeConfigToV3(); + + expect(configStore.get('userConfig')?.data).toEqual({}); + }); + + it('is a no-op when general is malformed', async () => { + configStore.set('userConfig', { + name: 'userConfig', + data: { + general: 'invalid', + }, + lastModified: Date.now(), + }); + + await upgradeConfigToV3(); + + expect(configStore.get('userConfig')?.data).toEqual({ + general: 'invalid', + }); + }); +}); diff --git a/src/core/config-upgrader/upgradeConfigToV3.ts b/src/core/config-upgrader/upgradeConfigToV3.ts new file mode 100644 index 0000000..22ffc63 --- /dev/null +++ b/src/core/config-upgrader/upgradeConfigToV3.ts @@ -0,0 +1,35 @@ +import { KGConfigStorage } from '../io/KGConfigStorage'; +import { LOCAL_LLM_DEFAULT_CONTEXT_LENGTH } from '../../util/localLLMConfig'; + +const CONFIG_KEY = 'userConfig'; + +export async function upgradeConfigToV3(): Promise { + const storage = KGConfigStorage.getInstance(); + const rawConfig = await storage.getRaw(CONFIG_KEY); + if (!rawConfig || typeof rawConfig !== 'object') { + return; + } + + const config = rawConfig as Record; + const general = config.general; + if (!general || typeof general !== 'object') { + return; + } + + const generalRecord = general as Record; + const localBrowser = generalRecord.local_browser; + + if (!localBrowser || typeof localBrowser !== 'object') { + generalRecord.local_browser = { context_length: LOCAL_LLM_DEFAULT_CONTEXT_LENGTH }; + await storage.saveRaw(CONFIG_KEY, config); + return; + } + + const localBrowserRecord = localBrowser as Record; + if ('context_length' in localBrowserRecord) { + return; + } + + localBrowserRecord.context_length = LOCAL_LLM_DEFAULT_CONTEXT_LENGTH; + await storage.saveRaw(CONFIG_KEY, config); +} diff --git a/src/core/config/ConfigManager.ts b/src/core/config/ConfigManager.ts index 118e048..88d50b2 100644 --- a/src/core/config/ConfigManager.ts +++ b/src/core/config/ConfigManager.ts @@ -8,6 +8,9 @@ interface AppConfig { language: string; llm_provider: 'local_browser' | 'openai' | 'gemini' | 'claude' | 'claude_openrouter' | 'openai_compatible'; persist_api_keys_non_localhost: boolean; + local_browser: { + context_length: 32768 | 65536 | 131072; + }; openai: { api_key: string; flex: boolean; @@ -208,6 +211,9 @@ export class ConfigManager { base_url: '', model: '' }, + local_browser: { + context_length: 32768 + }, soundfont: { base_url: 'https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/' }, diff --git a/src/util/localLLMConfig.test.ts b/src/util/localLLMConfig.test.ts new file mode 100644 index 0000000..0b6522b --- /dev/null +++ b/src/util/localLLMConfig.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from 'vitest'; +import { + formatLocalLLMContextLength, + LOCAL_LLM_DEFAULT_CONTEXT_LENGTH, + normalizeLocalLLMContextLength, +} from './localLLMConfig'; + +describe('localLLMConfig', () => { + it('defaults invalid context lengths to 32768', () => { + expect(normalizeLocalLLMContextLength(undefined)).toBe(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH); + expect(normalizeLocalLLMContextLength(12345)).toBe(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH); + }); + + it('formats context lengths using k suffixes', () => { + expect(formatLocalLLMContextLength(32768)).toBe('32k'); + expect(formatLocalLLMContextLength(65536)).toBe('64k'); + expect(formatLocalLLMContextLength(131072)).toBe('128k'); + }); +}); diff --git a/src/util/localLLMConfig.ts b/src/util/localLLMConfig.ts index e2f632d..4b30060 100644 --- a/src/util/localLLMConfig.ts +++ b/src/util/localLLMConfig.ts @@ -6,6 +6,10 @@ export const LOCAL_LLM_DISPLAY_NAME = 'Gemma 4 E4B'; export const LOCAL_LLM_LEGACY_FILENAMES = [ 'gemma-3n-E4B-it-int4-Web.litertlm', ]; +export const LOCAL_LLM_CONTEXT_LENGTH_OPTIONS = [32768, 65536, 131072] as const; +export const LOCAL_LLM_DEFAULT_CONTEXT_LENGTH = 32768; + +export type LocalLLMContextLength = typeof LOCAL_LLM_CONTEXT_LENGTH_OPTIONS[number]; export interface LocalLLMRuntimeSupport { supported: boolean; @@ -40,3 +44,16 @@ export function detectLocalLLMRuntimeSupport(): LocalLLMRuntimeSupport { reason, }; } + +export function isLocalLLMContextLength(value: unknown): value is LocalLLMContextLength { + return typeof value === 'number' + && (LOCAL_LLM_CONTEXT_LENGTH_OPTIONS as readonly number[]).includes(value); +} + +export function normalizeLocalLLMContextLength(value: unknown): LocalLLMContextLength { + return isLocalLLMContextLength(value) ? value : LOCAL_LLM_DEFAULT_CONTEXT_LENGTH; +} + +export function formatLocalLLMContextLength(value: LocalLLMContextLength): string { + return `${Math.round(value / 1024)}k`; +} From 1a051fea54a2976ec4d258dc2ef3b57511b8dbad Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 21:24:03 -0700 Subject: [PATCH 11/19] feat: added a welcome msg for local LLM users --- public/chat/welcome_local_llm.md | 21 ++++ .../messageFilter/UserMessageFilter.test.ts | 109 ++++++++++++++++++ src/util/messageFilter/UserMessageFilter.ts | 53 +++++++-- 3 files changed, 175 insertions(+), 8 deletions(-) create mode 100644 public/chat/welcome_local_llm.md create mode 100644 src/util/messageFilter/UserMessageFilter.test.ts diff --git a/public/chat/welcome_local_llm.md b/public/chat/welcome_local_llm.md new file mode 100644 index 0000000..724767a --- /dev/null +++ b/public/chat/welcome_local_llm.md @@ -0,0 +1,21 @@ +## Local LLM Mode + +Welcome to **K.G.Studio Musician Assistant** in local LLM mode. + +- No external API calls are required. Everything runs directly in your browser, with no extra API cost. +- This mode uses **Gemma 4 E4B** through **LiteRT-LM** with **WebGPU** acceleration. +- Recommended hardware: a GPU with at least **8 GB VRAM** or a system with at least **16 GB unified RAM**. +- Performance is more limited than larger cloud-hosted models, especially on harder planning, editing, and multi-step tasks. + +### Recommended Workflow +- Keep requests small and focused. +- Guide the model step by step toward the final goal. +- Work on smaller music regions instead of large full-song edits. +- Prefer simpler music arrangements when possible. +- Start a new conversation for each standalone task. + +### Use an External LLM Instead +- If you want a larger cloud or self-hosted model, open **Settings -> General -> LLM Provider** and switch away from **Local LLM (Browser)**. +- For a cloud model, you can use **OpenAI**, or choose **OpenAI Compatible** and enter a provider such as OpenRouter. +- For a self-hosted model, choose **OpenAI Compatible** and enter your server's **Base URL** and **Model**. +- After switching providers, start a new conversation so the chat uses the new model cleanly. diff --git a/src/util/messageFilter/UserMessageFilter.test.ts b/src/util/messageFilter/UserMessageFilter.test.ts new file mode 100644 index 0000000..8c50965 --- /dev/null +++ b/src/util/messageFilter/UserMessageFilter.test.ts @@ -0,0 +1,109 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { processUserMessage } from './UserMessageFilter'; +import { addWelcomeMessage } from '../../utils/chatMessageUtils'; + +const configState = new Map(); + +const configManagerMock = { + getIsInitialized: vi.fn(() => true), + initialize: vi.fn().mockResolvedValue(undefined), + get: vi.fn((key: string) => configState.get(key)), +}; + +vi.mock('../../core/config/ConfigManager', () => ({ + ConfigManager: { + instance: () => configManagerMock, + }, +})); + +vi.mock('../chatUtil', () => ({ + clearChatHistoryAndUI: vi.fn(), +})); + +vi.mock('../../stores/projectStore', () => ({ + useProjectStore: { + getState: () => ({ + setStatus: vi.fn(), + activeRegionId: null, + selectedRegionIds: [], + }), + }, +})); + +vi.mock('../../agent/core/SystemPrompts', () => ({ + SystemPrompts: { + getPromptWithContext: vi.fn(async (value: string) => value), + }, +})); + +describe('processUserMessage /welcome', () => { + beforeEach(() => { + configState.clear(); + configState.set('general.llm_provider', 'local_browser'); + configState.set('general.openai.api_key', ''); + configState.set('general.gemini.api_key', ''); + configState.set('general.claude.api_key', ''); + configState.set('general.claude_openrouter.api_key', ''); + configState.set('general.openai_compatible.base_url', ''); + configState.set('general.openai_compatible.model', ''); + + configManagerMock.getIsInitialized.mockReturnValue(true); + configManagerMock.initialize.mockClear(); + configManagerMock.get.mockClear(); + + vi.stubGlobal('fetch', vi.fn(async (input: string | URL | Request) => { + const url = String(input); + return { + ok: true, + status: 200, + text: async () => `content:${url}`, + }; + })); + }); + + it('uses the local welcome for the local browser provider', async () => { + configState.set('general.llm_provider', 'local_browser'); + configState.set('general.openai.api_key', ''); + configState.set('general.openai_compatible.base_url', ''); + configState.set('general.openai_compatible.model', ''); + + const result = await processUserMessage('/welcome'); + + expect(fetch).toHaveBeenCalledWith(expect.stringContaining('chat/welcome_local_llm.md')); + expect(result.metadata).toMatchObject({ command: 'welcome', variant: 'local' }); + expect(result.pseudoAssistantResponse).toContain('welcome_local_llm.md'); + }); + + it('uses the new-user welcome for non-local providers without required config', async () => { + configState.set('general.llm_provider', 'openai'); + configState.set('general.openai.api_key', ''); + + const result = await processUserMessage('/welcome'); + + expect(fetch).toHaveBeenCalledWith(expect.stringContaining('chat/welcome_new.md')); + expect(result.metadata).toMatchObject({ command: 'welcome', variant: 'new' }); + expect(result.pseudoAssistantResponse).toContain('welcome_new.md'); + }); + + it('uses the returning-user welcome for configured non-local providers', async () => { + configState.set('general.llm_provider', 'openai_compatible'); + configState.set('general.openai_compatible.base_url', 'https://openrouter.ai/api/v1'); + configState.set('general.openai_compatible.model', 'qwen/qwen3-30b-a3b:free'); + + const result = await processUserMessage('/welcome'); + + expect(fetch).toHaveBeenCalledWith(expect.stringContaining('chat/welcome_again.md')); + expect(result.metadata).toMatchObject({ command: 'welcome', variant: 'again' }); + expect(result.pseudoAssistantResponse).toContain('welcome_again.md'); + }); + + it('reuses the same welcome routing through addWelcomeMessage', async () => { + configState.set('general.llm_provider', 'local_browser'); + + const message = await addWelcomeMessage(); + + expect(fetch).toHaveBeenCalledWith(expect.stringContaining('chat/welcome_local_llm.md')); + expect(message?.role).toBe('assistant'); + expect(message?.content).toContain('welcome_local_llm.md'); + }); +}); diff --git a/src/util/messageFilter/UserMessageFilter.ts b/src/util/messageFilter/UserMessageFilter.ts index b314b18..b168f85 100644 --- a/src/util/messageFilter/UserMessageFilter.ts +++ b/src/util/messageFilter/UserMessageFilter.ts @@ -17,6 +17,48 @@ export interface UserMessageFilterResult { metadata?: Record; } +function hasText(value: unknown): boolean { + return typeof value === 'string' && value.trim() !== ''; +} + +function getWelcomeVariant(configManager: ConfigManager): 'local' | 'new' | 'again' { + const provider = (configManager.get('general.llm_provider') as string) || LOCAL_LLM_PROVIDER_KEY; + + if (provider === LOCAL_LLM_PROVIDER_KEY) { + return 'local'; + } + + switch (provider) { + case 'openai': + return hasText(configManager.get('general.openai.api_key')) ? 'again' : 'new'; + case 'gemini': + return hasText(configManager.get('general.gemini.api_key')) ? 'again' : 'new'; + case 'claude': + return hasText(configManager.get('general.claude.api_key')) ? 'again' : 'new'; + case 'claude_openrouter': + return hasText(configManager.get('general.claude_openrouter.api_key')) ? 'again' : 'new'; + case 'openai_compatible': + return hasText(configManager.get('general.openai_compatible.base_url')) + && hasText(configManager.get('general.openai_compatible.model')) + ? 'again' + : 'new'; + default: + return 'new'; + } +} + +function getWelcomeUrl(variant: 'local' | 'new' | 'again'): string { + switch (variant) { + case 'local': + return `${import.meta.env.BASE_URL}chat/welcome_local_llm.md`; + case 'again': + return `${import.meta.env.BASE_URL}chat/welcome_again.md`; + case 'new': + default: + return `${import.meta.env.BASE_URL}chat/welcome_new.md`; + } +} + /** * Process a user message before it is displayed or sent to the LLM. * Handles slash-commands and returns a structured decision. @@ -53,12 +95,8 @@ export async function processUserMessage(originalMessage: string): Promise Date: Thu, 14 May 2026 21:37:06 -0700 Subject: [PATCH 12/19] fix: switch to processing status when receiving response tokens from LLM --- src/components/chat/AssistantMessage.test.tsx | 25 +++++ src/components/chat/AssistantMessage.tsx | 12 ++- src/hooks/useStreamProcessor.test.ts | 92 +++++++++++++++++++ src/hooks/useStreamProcessor.ts | 4 +- src/test/setup.ts | 14 ++- 5 files changed, 139 insertions(+), 8 deletions(-) create mode 100644 src/components/chat/AssistantMessage.test.tsx create mode 100644 src/hooks/useStreamProcessor.test.ts diff --git a/src/components/chat/AssistantMessage.test.tsx b/src/components/chat/AssistantMessage.test.tsx new file mode 100644 index 0000000..04d68b2 --- /dev/null +++ b/src/components/chat/AssistantMessage.test.tsx @@ -0,0 +1,25 @@ +import { fireEvent, render, screen } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; +import AssistantMessage from './AssistantMessage'; + +describe('AssistantMessage', () => { + it.each([ + 'Thinking... click here to abort.', + 'Processing... 3 tokens received. click here to abort.' + ])('renders the abort action for streaming status content: %s', (content) => { + const onAbort = vi.fn(); + + render( + + ); + + const abortButton = screen.getByRole('button', { name: 'click here to abort' }); + expect(abortButton).toBeInTheDocument(); + fireEvent.click(abortButton); + expect(onAbort).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx index ad196be..96b2b93 100644 --- a/src/components/chat/AssistantMessage.tsx +++ b/src/components/chat/AssistantMessage.tsx @@ -44,22 +44,26 @@ const AssistantMessage: React.FC = ({ content, isStreamin const prefillTps = formatTps(performanceInfo?.prefillTps); const generationTps = formatTps(performanceInfo?.generationTps); const hasPerformanceInfo = Boolean(prefillTps || generationTps); + const processingWaveLabels = ['Thinking...', 'Processing...']; const renderContent = () => { // Handle special abort link for streaming messages if (isStreaming && onAbort && content.includes('click here to abort')) { - const hasProcessingWave = content.includes('Thinking...'); + const processingWaveMarkup = processingWaveLabels + .map(label => `${label}`) + .find(markup => content.includes(markup)); - if (hasProcessingWave) { + if (processingWaveMarkup) { const parts = content.split('click here to abort'); const beforeAbort = parts[0].replace( - 'Thinking...', + processingWaveMarkup, '' ); + const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...'; return ( - Thinking... + {waveLabel} {beforeAbort}
{activeProvider === LOCAL_LLM_PROVIDER_KEY && ( -
-
-

{LOCAL_LLM_DISPLAY_NAME} Local Runtime

+
+
+

{LOCAL_LLM_DISPLAY_NAME} Local Runtime

{!localModelState.runtimeSupport.supported && ( -
+
{localModelState.runtimeSupport.reason}
)} {!localModelState.isCached && !localModelState.isDownloading && localModelState.runtimeSupport.supported && ( -
+
The local language model has not been downloaded yet. It will be downloaded automatically the next time you send a chat request with this provider.
)} {(localModelState.isChecking || localModelState.isDownloading || localModelState.progressText) && ( -
+
-
+
-
+
{localModelState.isChecking ? 'Checking local model cache...' : localModelState.progressText}
)} {localModelState.error && ( -
+
{localModelState.error}
)} From f9a3edb832ff4e77fae496b2e0cba06f645fcb4a Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 21:55:59 -0700 Subject: [PATCH 14/19] fix: update default playhead fps to 30 --- public/config.json | 2 +- src/components/settings/Settings.css | 10 +++++----- src/components/settings/sections/BehaviorSettings.tsx | 4 ++-- src/core/KGCore.ts | 2 +- src/core/config/ConfigManager.ts | 2 +- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/public/config.json b/public/config.json index 015918a..88226c1 100644 --- a/public/config.json +++ b/public/config.json @@ -69,7 +69,7 @@ } }, "editor": { - "playhead_update_frequency": 10, + "playhead_update_frequency": 30, "spectrogram_height_resolution": 3 }, "chatbox": { diff --git a/src/components/settings/Settings.css b/src/components/settings/Settings.css index bf85606..03c66ce 100644 --- a/src/components/settings/Settings.css +++ b/src/components/settings/Settings.css @@ -299,17 +299,17 @@ .settings-progress-track { width: 100%; - height: 10px; + height: 8px; border-radius: 999px; overflow: hidden; - background-color: #3a3a3a; - border: 1px solid #4a4a4a; + background-color: #1d1d1d; + border: 1px solid #3a3a3a; } .settings-progress-fill { height: 100%; - background: linear-gradient(90deg, #5a9fd4 0%, #76c28f 100%); - transition: width 0.2s ease; + background: linear-gradient(90deg, #5a9fd4 0%, #7cc2f1 100%); + transition: width 0.15s linear; } /* Settings Help Links */ diff --git a/src/components/settings/sections/BehaviorSettings.tsx b/src/components/settings/sections/BehaviorSettings.tsx index 892a30b..88f6d05 100644 --- a/src/components/settings/sections/BehaviorSettings.tsx +++ b/src/components/settings/sections/BehaviorSettings.tsx @@ -7,7 +7,7 @@ import { } from '../../../util/spectrogramUtil'; const BehaviorSettings: React.FC = () => { - const [playheadUpdateFrequency, setPlayheadUpdateFrequency] = useState(10); + const [playheadUpdateFrequency, setPlayheadUpdateFrequency] = useState(30); const [spectrogramHeightResolution, setSpectrogramHeightResolution] = useState(3); const [chatboxDefaultOpen, setChatboxDefaultOpen] = useState(true); const [audioLookaheadTime, setAudioLookaheadTime] = useState('50'); @@ -28,7 +28,7 @@ const BehaviorSettings: React.FC = () => { await configManager.initialize(); } - setPlayheadUpdateFrequency((configManager.get('editor.playhead_update_frequency') as number) ?? 10); + setPlayheadUpdateFrequency((configManager.get('editor.playhead_update_frequency') as number) ?? 30); setSpectrogramHeightResolution( normalizeSpectrogramHeightResolution(configManager.get('editor.spectrogram_height_resolution')) ); diff --git a/src/core/KGCore.ts b/src/core/KGCore.ts index 43eabf0..06dbfe8 100644 --- a/src/core/KGCore.ts +++ b/src/core/KGCore.ts @@ -367,7 +367,7 @@ export class KGCore { // Get playhead update frequency from config (in fps) const configManager = ConfigManager.instance(); - const updateFrequency = (configManager.get('editor.playhead_update_frequency') as number) ?? 10; + const updateFrequency = (configManager.get('editor.playhead_update_frequency') as number) ?? 30; const updateIntervalMs = 1000 / updateFrequency; // Convert fps to milliseconds this.playbackIntervalId = window.setInterval(() => { diff --git a/src/core/config/ConfigManager.ts b/src/core/config/ConfigManager.ts index 88d50b2..0b6a102 100644 --- a/src/core/config/ConfigManager.ts +++ b/src/core/config/ConfigManager.ts @@ -253,7 +253,7 @@ export class ConfigManager { }, }, editor: { - playhead_update_frequency: 10, + playhead_update_frequency: 30, spectrogram_height_resolution: 3 }, chatbox: { From 892e154a0f7867f2991f1ba30175e49a61b84219 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 22:35:04 -0700 Subject: [PATCH 15/19] fix: minor UI adjustments --- src/components/KGOnePanel.css | 25 +++-- src/components/KGOnePanel.tsx | 16 ++-- src/components/Toolbar.test.tsx | 165 ++++++++++++++++++++++++++++++++ src/components/Toolbar.tsx | 28 ++++-- src/index.css | 5 + src/stores/projectStore.test.ts | 76 +++++++++++++++ src/stores/projectStore.ts | 84 +++++++++++++--- 7 files changed, 368 insertions(+), 31 deletions(-) create mode 100644 src/components/Toolbar.test.tsx diff --git a/src/components/KGOnePanel.css b/src/components/KGOnePanel.css index d7d79c9..dbfc3bf 100644 --- a/src/components/KGOnePanel.css +++ b/src/components/KGOnePanel.css @@ -296,24 +296,37 @@ /* Generate / Separate button */ .kgone-btn-generate { - background-color: #4a5fa0; - border: none; border-radius: 4px; - color: #e0e0e0; font-size: 12px; font-weight: 600; padding: 8px 12px; - cursor: pointer; width: 100%; margin-top: 6px; +} + +.kgone-btn-generate-accent { + background-color: #4a5fa0; + border: none; + color: #e0e0e0; + cursor: pointer; transition: background-color 0.15s; } -.kgone-btn-generate:hover:not(:disabled) { +.kgone-btn-generate-accent:hover:not(:disabled) { background-color: #5a70b8; } -.kgone-btn-generate:disabled { +.kgone-btn-generate-accent:disabled { + background-color: #3a3a3a; + color: #666; + cursor: not-allowed; +} + +.kgone-btn-generate.dialog-btn { + min-height: 32px; +} + +.kgone-btn-generate.dialog-btn:disabled { background-color: #3a3a3a; color: #666; cursor: not-allowed; diff --git a/src/components/KGOnePanel.tsx b/src/components/KGOnePanel.tsx index a0d48ce..4c26bdb 100644 --- a/src/components/KGOnePanel.tsx +++ b/src/components/KGOnePanel.tsx @@ -524,7 +524,7 @@ const ClipTab: React.FC = ({ bpm, keySignature }) => { )} ) : ( - + )} diff --git a/src/index.css b/src/index.css index 4bb29a6..74651de 100644 --- a/src/index.css +++ b/src/index.css @@ -33,6 +33,11 @@ body { min-height: 100vh; } +a, +a:visited { + color: rgb(90, 159, 212); +} + #root { width: 100%; height: 100%; diff --git a/src/stores/projectStore.test.ts b/src/stores/projectStore.test.ts index 5457b24..66b1478 100644 --- a/src/stores/projectStore.test.ts +++ b/src/stores/projectStore.test.ts @@ -250,4 +250,80 @@ describe('projectStore piano roll state', () => { expect(useProjectStore.getState().trackAutomationRedrawVersion).toBe(initialVersion + 2); }); + + it('restores Chat after closing Settings when Chat was active on entry', async () => { + const { useProjectStore } = await import('./projectStore'); + + act(() => { + useProjectStore.getState().toggleChatBox(); + useProjectStore.getState().setShowSettings(true); + }); + + let state = useProjectStore.getState(); + expect(state.showSettings).toBe(true); + expect(state.settingsReturnSidePanel).toBe('chat'); + + act(() => { + useProjectStore.getState().setShowSettings(false); + }); + + state = useProjectStore.getState(); + expect(state.showSettings).toBe(false); + expect(state.showChatBox).toBe(true); + expect(state.showKGOnePanel).toBe(false); + expect(state.showEventListPanel).toBe(false); + }); + + it('restores K.G.One after closing Settings when K.G.One was active on entry', async () => { + const { useProjectStore } = await import('./projectStore'); + + act(() => { + useProjectStore.getState().toggleKGOnePanel(); + useProjectStore.getState().setShowSettings(true); + useProjectStore.getState().setShowSettings(false); + }); + + const state = useProjectStore.getState(); + expect(state.showSettings).toBe(false); + expect(state.showKGOnePanel).toBe(true); + expect(state.showChatBox).toBe(false); + expect(state.showEventListPanel).toBe(false); + }); + + it('restores no side panel after closing Settings when none was active on entry', async () => { + const { useProjectStore } = await import('./projectStore'); + + act(() => { + useProjectStore.getState().toggleChatBox(); + useProjectStore.getState().toggleChatBox(); + useProjectStore.getState().setShowSettings(true); + useProjectStore.getState().setShowSettings(false); + }); + + const state = useProjectStore.getState(); + expect(state.showSettings).toBe(false); + expect(state.showChatBox).toBe(false); + expect(state.showKGOnePanel).toBe(false); + expect(state.showEventListPanel).toBe(false); + expect(state.settingsReturnSidePanel).toBeNull(); + expect(state.lastActiveSidePanel).toBe('chat'); + }); + + it('opens Event List and exits Settings when Event List is activated from Settings', async () => { + const { useProjectStore } = await import('./projectStore'); + + act(() => { + useProjectStore.getState().toggleChatBox(); + useProjectStore.getState().setShowSettings(true); + useProjectStore.getState().activateSidePanel('eventList'); + }); + + const state = useProjectStore.getState(); + expect(state.showSettings).toBe(false); + expect(state.showEventListPanel).toBe(true); + expect(state.showChatBox).toBe(false); + expect(state.showKGOnePanel).toBe(false); + expect(state.settingsReturnSidePanel).toBeNull(); + expect(state.lastActiveSidePanel).toBe('eventList'); + }); }); diff --git a/src/stores/projectStore.ts b/src/stores/projectStore.ts index 32f34cf..c8bf4dd 100644 --- a/src/stores/projectStore.ts +++ b/src/stores/projectStore.ts @@ -53,6 +53,16 @@ function updateBarWidthMultiplierCSS(multiplier: number): void { ); } +type SidePanelType = 'kgone' | 'chat' | 'eventList'; + +function getSidePanelVisibilityState(activePanel: SidePanelType | null) { + return { + showKGOnePanel: activePanel === 'kgone', + showChatBox: activePanel === 'chat', + showEventListPanel: activePanel === 'eventList', + }; +} + // Define the store state interface interface ProjectState { // State @@ -101,6 +111,8 @@ interface ProjectState { // Event list panel state showEventListPanel: boolean; + lastActiveSidePanel: SidePanelType | null; + settingsReturnSidePanel: SidePanelType | null; // Instrument selection panel state showInstrumentSelection: boolean; @@ -199,6 +211,7 @@ interface ProjectState { // Event List panel actions toggleEventListPanel: () => void; + activateSidePanel: (panel: SidePanelType) => void; // Instrument selection panel actions openInstrumentSelectionForTrack: () => void; @@ -430,6 +443,8 @@ export const useProjectStore = create((set, get) => { // Initial Event List panel state showEventListPanel: false, + lastActiveSidePanel: initialChatBoxState ? 'chat' : null, + settingsReturnSidePanel: null, // Initial Instrument Selection panel state showInstrumentSelection: initialShowInstrumentSelection, @@ -1579,27 +1594,52 @@ export const useProjectStore = create((set, get) => { }, // ChatBox action implementations - setShowChatBox: (show: boolean) => { + activateSidePanel: (panel: SidePanelType) => { set({ - showChatBox: show, - showKGOnePanel: show ? false : get().showKGOnePanel, - showEventListPanel: show ? false : get().showEventListPanel + ...getSidePanelVisibilityState(panel), + lastActiveSidePanel: panel, + showSettings: false, + settingsReturnSidePanel: null, }); }, + setShowChatBox: (show: boolean) => { + if (show) { + get().activateSidePanel('chat'); + return; + } + + set({ showChatBox: false }); + }, + toggleChatBox: () => { - const { showChatBox } = get(); - set({ showChatBox: !showChatBox, showKGOnePanel: false, showEventListPanel: false }); + const { showChatBox, showSettings } = get(); + if (showSettings || !showChatBox) { + get().activateSidePanel('chat'); + return; + } + + set({ showChatBox: false }); }, toggleKGOnePanel: () => { - const { showKGOnePanel } = get(); - set({ showKGOnePanel: !showKGOnePanel, showChatBox: false, showEventListPanel: false }); + const { showKGOnePanel, showSettings } = get(); + if (showSettings || !showKGOnePanel) { + get().activateSidePanel('kgone'); + return; + } + + set({ showKGOnePanel: false }); }, toggleEventListPanel: () => { - const { showEventListPanel } = get(); - set({ showEventListPanel: !showEventListPanel, showChatBox: false, showKGOnePanel: false }); + const { showEventListPanel, showSettings } = get(); + if (showSettings || !showEventListPanel) { + get().activateSidePanel('eventList'); + return; + } + + set({ showEventListPanel: false }); }, // Instrument selection panel actions @@ -1615,7 +1655,29 @@ export const useProjectStore = create((set, get) => { // Settings action implementations setShowSettings: (show: boolean) => { - set({ showSettings: show }); + if (show) { + const { showKGOnePanel, showChatBox, showEventListPanel } = get(); + const activePanel = showKGOnePanel + ? 'kgone' + : showChatBox + ? 'chat' + : showEventListPanel + ? 'eventList' + : null; + + set({ + showSettings: true, + settingsReturnSidePanel: activePanel, + }); + return; + } + + const { settingsReturnSidePanel } = get(); + set({ + showSettings: false, + settingsReturnSidePanel: null, + ...getSidePanelVisibilityState(settingsReturnSidePanel), + }); }, toggleSettings: () => { From 42b9662bb40f781941ea36308f76135f03685817 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 23:15:38 -0700 Subject: [PATCH 16/19] fix: minor UI adjustments --- src/components/piano-roll/PianoRoll.css | 14 +++++++------- src/util/scaleUtil.test.ts | 2 ++ src/util/scaleUtil.ts | 11 +++++------ 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/components/piano-roll/PianoRoll.css b/src/components/piano-roll/PianoRoll.css index 56730dd..56c8d46 100644 --- a/src/components/piano-roll/PianoRoll.css +++ b/src/components/piano-roll/PianoRoll.css @@ -610,9 +610,9 @@ position: absolute; left: 0; right: 0; - background-color: rgba(123, 104, 238, 0.15); - border-top: 1px solid rgba(123, 104, 238, 0.3); - border-bottom: 1px solid rgba(123, 104, 238, 0.3); + background-color: rgba(135, 206, 250, 0.15); + border-top: 1px solid rgba(135, 206, 250, 0.3); + border-bottom: 1px solid rgba(135, 206, 250, 0.3); pointer-events: none; z-index: 1; transition: opacity 0.1s ease; @@ -622,9 +622,9 @@ position: absolute; top: 0; bottom: 0; - background-color: rgba(123, 104, 238, 0.1); - border-left: 1px solid rgba(123, 104, 238, 0.25); - border-right: 1px solid rgba(123, 104, 238, 0.25); + background-color: rgba(135, 206, 250, 0.1); + border-left: 1px solid rgba(135, 206, 250, 0.25); + border-right: 1px solid rgba(135, 206, 250, 0.25); pointer-events: none; z-index: 1; transition: opacity 0.1s ease; @@ -802,4 +802,4 @@ color: #e0e0e0; min-width: 20px; text-align: center; -} \ No newline at end of file +} diff --git a/src/util/scaleUtil.test.ts b/src/util/scaleUtil.test.ts index baa6852..c1c11a0 100644 --- a/src/util/scaleUtil.test.ts +++ b/src/util/scaleUtil.test.ts @@ -385,6 +385,8 @@ describe('scaleUtil', () => { const result = generatePianoGridBackground('ionian', 'C major'); expect(result).toContain('linear-gradient'); expect(typeof result).toBe('string'); + expect(result).toContain('#282828'); + expect(result).toContain('#303030'); }); it('should generate different backgrounds for different modes', () => { diff --git a/src/util/scaleUtil.ts b/src/util/scaleUtil.ts index 0b87037..5cb72ac 100644 --- a/src/util/scaleUtil.ts +++ b/src/util/scaleUtil.ts @@ -320,19 +320,18 @@ export const generatePianoGridBackground = ( const rowBottomMinusOne = `calc(var(--region-piano-key-height) * ${index + 1} - 1px)`; const rowBottom = `calc(var(--region-piano-key-height) * ${index + 1})`; - // For scale notes: highlight the full row with a semi-transparent blue background - // For non-scale notes: use transparent background with just the separator line + // Match the event list palette while preserving scale-aware row distinction. if (isInScale) { return ` - rgba(90, 123, 154, 0.15) ${rowTop}, - rgba(90, 123, 154, 0.15) ${rowBottomMinusOne}, + #282828 ${rowTop}, + #282828 ${rowBottomMinusOne}, #3a3a3a ${rowBottomMinusOne}, #3a3a3a ${rowBottom} `.trim(); } else { return ` - transparent ${rowTop}, - transparent ${rowBottomMinusOne}, + #303030 ${rowTop}, + #303030 ${rowBottomMinusOne}, #3a3a3a ${rowBottomMinusOne}, #3a3a3a ${rowBottom} `.trim(); From ac2d9247117daac1d275051b4649ce48e0e73e4f Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Fri, 15 May 2026 12:07:45 -0700 Subject: [PATCH 17/19] feat: added thinking time statistics; added LaTeX rendering support --- package-lock.json | 298 +++++++++++++++++- package.json | 3 + src/components/ChatBox.css | 10 + src/components/chat/AssistantMessage.test.tsx | 75 ++++- src/components/chat/AssistantMessage.tsx | 47 ++- src/main.tsx | 1 + 6 files changed, 426 insertions(+), 8 deletions(-) diff --git a/package-lock.json b/package-lock.json index bc951f5..7a0f770 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,6 +13,7 @@ "fft.js": "^4.0.4", "idb": "^8.0.3", "jszip": "^3.10.1", + "katex": "^0.16.46", "onnxruntime-web": "^1.26.0", "openai": "^6.33.0", "react": "^19.1.0", @@ -21,7 +22,9 @@ "react-markdown": "^10.1.0", "react-syntax-highlighter": "^15.6.1", "reflect-metadata": "^0.2.2", + "rehype-katex": "^7.0.1", "remark-gfm": "^4.0.1", + "remark-math": "^6.0.0", "tone": "^15.1.22", "vexflow": "^5.0.0", "zustand": "^5.0.6" @@ -2669,6 +2672,11 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/katex": { + "version": "0.16.8", + "resolved": "https://registry.npmjs.org/@types/katex/-/katex-0.16.8.tgz", + "integrity": "sha512-trgaNyfU+Xh2Tc+ABIb44a5AYUpicB3uwirOioeOkNPPbmgRNtcWyDeeFRzjPZENO9Vq8gvVqfhaaXWLlevVwg==" + }, "node_modules/@types/mdast": { "version": "4.0.4", "resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-4.0.4.tgz", @@ -3745,6 +3753,14 @@ "url": "https://github.com/sponsors/wooorm" } }, + "node_modules/commander": { + "version": "8.3.0", + "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", + "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", + "engines": { + "node": ">= 12" + } + }, "node_modules/compare-func": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/compare-func/-/compare-func-2.0.0.tgz", @@ -4202,7 +4218,6 @@ "version": "6.0.1", "resolved": "https://registry.npmjs.org/entities/-/entities-6.0.1.tgz", "integrity": "sha512-aN97NXWF6AWBTahfVOIrB/NShkzi5H7F9r1s9mD3cDj4Ko5f2qhhVoYMibXF7GlLveb/D2ioWay8lxI97Ven3g==", - "dev": true, "license": "BSD-2-Clause", "engines": { "node": ">=0.12" @@ -5218,6 +5233,139 @@ "node": ">= 0.4" } }, + "node_modules/hast-util-from-dom": { + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/hast-util-from-dom/-/hast-util-from-dom-5.0.1.tgz", + "integrity": "sha512-N+LqofjR2zuzTjCPzyDUdSshy4Ma6li7p/c3pA78uTwzFgENbgbUrm2ugwsOdcjI1muO+o6Dgzp9p8WHtn/39Q==", + "dependencies": { + "@types/hast": "^3.0.0", + "hastscript": "^9.0.0", + "web-namespaces": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-dom/node_modules/hast-util-parse-selector": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/hast-util-parse-selector/-/hast-util-parse-selector-4.0.0.tgz", + "integrity": "sha512-wkQCkSYoOGCRKERFWcxMVMOcYE2K1AaNLU8DXS9arxnLOUEWbOXKXiJUNzEpqZ3JOKpnha3jkFrumEjVliDe7A==", + "dependencies": { + "@types/hast": "^3.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-dom/node_modules/hastscript": { + "version": "9.0.1", + "resolved": "https://registry.npmjs.org/hastscript/-/hastscript-9.0.1.tgz", + "integrity": "sha512-g7df9rMFX/SPi34tyGCyUBREQoKkapwdY/T04Qn9TDWfHhAYt4/I0gMVirzK5wEzeUqIjEB+LXC/ypb7Aqno5w==", + "dependencies": { + "@types/hast": "^3.0.0", + "comma-separated-tokens": "^2.0.0", + "hast-util-parse-selector": "^4.0.0", + "property-information": "^7.0.0", + "space-separated-tokens": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-html": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/hast-util-from-html/-/hast-util-from-html-2.0.3.tgz", + "integrity": "sha512-CUSRHXyKjzHov8yKsQjGOElXy/3EKpyX56ELnkHH34vDVw1N1XSQ1ZcAvTyAPtGqLTuKP/uxM+aLkSPqF/EtMw==", + "dependencies": { + "@types/hast": "^3.0.0", + "devlop": "^1.1.0", + "hast-util-from-parse5": "^8.0.0", + "parse5": "^7.0.0", + "vfile": "^6.0.0", + "vfile-message": "^4.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-html-isomorphic": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/hast-util-from-html-isomorphic/-/hast-util-from-html-isomorphic-2.0.0.tgz", + "integrity": "sha512-zJfpXq44yff2hmE0XmwEOzdWin5xwH+QIhMLOScpX91e/NSGPsAzNCvLQDIEPyO2TXi+lBmU6hjLIhV8MwP2kw==", + "dependencies": { + "@types/hast": "^3.0.0", + "hast-util-from-dom": "^5.0.0", + "hast-util-from-html": "^2.0.0", + "unist-util-remove-position": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-parse5": { + "version": "8.0.3", + "resolved": "https://registry.npmjs.org/hast-util-from-parse5/-/hast-util-from-parse5-8.0.3.tgz", + "integrity": "sha512-3kxEVkEKt0zvcZ3hCRYI8rqrgwtlIOFMWkbclACvjlDw8Li9S2hk/d51OI0nr/gIpdMHNepwgOKqZ/sy0Clpyg==", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/unist": "^3.0.0", + "devlop": "^1.0.0", + "hastscript": "^9.0.0", + "property-information": "^7.0.0", + "vfile": "^6.0.0", + "vfile-location": "^5.0.0", + "web-namespaces": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-parse5/node_modules/hast-util-parse-selector": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/hast-util-parse-selector/-/hast-util-parse-selector-4.0.0.tgz", + "integrity": "sha512-wkQCkSYoOGCRKERFWcxMVMOcYE2K1AaNLU8DXS9arxnLOUEWbOXKXiJUNzEpqZ3JOKpnha3jkFrumEjVliDe7A==", + "dependencies": { + "@types/hast": "^3.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-parse5/node_modules/hastscript": { + "version": "9.0.1", + "resolved": "https://registry.npmjs.org/hastscript/-/hastscript-9.0.1.tgz", + "integrity": "sha512-g7df9rMFX/SPi34tyGCyUBREQoKkapwdY/T04Qn9TDWfHhAYt4/I0gMVirzK5wEzeUqIjEB+LXC/ypb7Aqno5w==", + "dependencies": { + "@types/hast": "^3.0.0", + "comma-separated-tokens": "^2.0.0", + "hast-util-parse-selector": "^4.0.0", + "property-information": "^7.0.0", + "space-separated-tokens": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-is-element": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/hast-util-is-element/-/hast-util-is-element-3.0.0.tgz", + "integrity": "sha512-Val9mnv2IWpLbNPqc/pUem+a7Ipj2aHacCwgNfTiK0vJKl0LF+4Ba4+v1oPHFpf3bLYmreq0/l3Gud9S5OH42g==", + "dependencies": { + "@types/hast": "^3.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/hast-util-parse-selector": { "version": "2.2.5", "resolved": "https://registry.npmjs.org/hast-util-parse-selector/-/hast-util-parse-selector-2.2.5.tgz", @@ -5255,6 +5403,21 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/hast-util-to-text": { + "version": "4.0.2", + "resolved": "https://registry.npmjs.org/hast-util-to-text/-/hast-util-to-text-4.0.2.tgz", + "integrity": "sha512-KK6y/BN8lbaq654j7JgBydev7wuNMcID54lkRav1P0CaE1e47P72AWWPiGKXTJU271ooYzcvTAn/Zt0REnvc7A==", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/unist": "^3.0.0", + "hast-util-is-element": "^3.0.0", + "unist-util-find-after": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/hast-util-whitespace": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/hast-util-whitespace/-/hast-util-whitespace-3.0.0.tgz", @@ -6038,6 +6201,21 @@ "setimmediate": "^1.0.5" } }, + "node_modules/katex": { + "version": "0.16.46", + "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.46.tgz", + "integrity": "sha512-WHy4Coo+bGZyH7NwJKHkS04YFsFcarWbAEOAC3EMndzdN6VSZqklLLIgfxzyaW9jDoeGYJX9SWbJPKpecox0Uw==", + "funding": [ + "https://opencollective.com/katex", + "https://github.com/sponsors/katex" + ], + "dependencies": { + "commander": "^8.3.0" + }, + "bin": { + "katex": "cli.js" + } + }, "node_modules/keyv": { "version": "4.5.4", "resolved": "https://registry.npmjs.org/keyv/-/keyv-4.5.4.tgz", @@ -6529,6 +6707,24 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/mdast-util-math": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/mdast-util-math/-/mdast-util-math-3.0.0.tgz", + "integrity": "sha512-Tl9GBNeG/AhJnQM221bJR2HPvLOSnLE/T9cJI9tlc6zwQk2nPk/4f0cHkOdEixQPC/j8UtKDdITswvLAy1OZ1w==", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/mdast": "^4.0.0", + "devlop": "^1.0.0", + "longest-streak": "^3.0.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.1.0", + "unist-util-remove-position": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/mdast-util-mdx-expression": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/mdast-util-mdx-expression/-/mdast-util-mdx-expression-2.0.1.tgz", @@ -7031,6 +7227,24 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/micromark-extension-math": { + "version": "3.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-math/-/micromark-extension-math-3.1.0.tgz", + "integrity": "sha512-lvEqd+fHjATVs+2v/8kg9i5Q0AP2k85H0WUOwpIVvUML8BapsMvh1XAogmQjOCsLpoKRCVQqEkQBB3NhVBcsOg==", + "dependencies": { + "@types/katex": "^0.16.0", + "devlop": "^1.0.0", + "katex": "^0.16.0", + "micromark-factory-space": "^2.0.0", + "micromark-util-character": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/micromark-factory-destination": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/micromark-factory-destination/-/micromark-factory-destination-2.0.1.tgz", @@ -11239,7 +11453,6 @@ "version": "7.3.0", "resolved": "https://registry.npmjs.org/parse5/-/parse5-7.3.0.tgz", "integrity": "sha512-IInvU7fabl34qmi9gY8XOVxhYyMyuH2xUNpb2q8/Y+7552KlejkRvqvD19nMoUW/uQGGbqNpA6Tufu5FL5BZgw==", - "dev": true, "license": "MIT", "dependencies": { "entities": "^6.0.0" @@ -12139,6 +12352,24 @@ "node": ">=14" } }, + "node_modules/rehype-katex": { + "version": "7.0.1", + "resolved": "https://registry.npmjs.org/rehype-katex/-/rehype-katex-7.0.1.tgz", + "integrity": "sha512-OiM2wrZ/wuhKkigASodFoo8wimG3H12LWQaH8qSPVJn9apWKFSH3YOCtbKpBorTVw/eI7cuT21XBbvwEswbIOA==", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/katex": "^0.16.0", + "hast-util-from-html-isomorphic": "^2.0.0", + "hast-util-to-text": "^4.0.0", + "katex": "^0.16.0", + "unist-util-visit-parents": "^6.0.0", + "vfile": "^6.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/remark-gfm": { "version": "4.0.1", "resolved": "https://registry.npmjs.org/remark-gfm/-/remark-gfm-4.0.1.tgz", @@ -12157,6 +12388,21 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/remark-math": { + "version": "6.0.0", + "resolved": "https://registry.npmjs.org/remark-math/-/remark-math-6.0.0.tgz", + "integrity": "sha512-MMqgnP74Igy+S3WwnhQ7kqGlEerTETXMvJhrUzDikVZ2/uogJCb+WHUg97hK9/jcfc0dkD73s3LN8zU49cTEtA==", + "dependencies": { + "@types/mdast": "^4.0.0", + "mdast-util-math": "^3.0.0", + "micromark-extension-math": "^3.0.0", + "unified": "^11.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/remark-parse": { "version": "11.0.0", "resolved": "https://registry.npmjs.org/remark-parse/-/remark-parse-11.0.0.tgz", @@ -13713,6 +13959,19 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/unist-util-find-after": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/unist-util-find-after/-/unist-util-find-after-5.0.0.tgz", + "integrity": "sha512-amQa0Ep2m6hE2g72AugUItjbuM8X8cGQnFoHk0pGfrFeT9GZhzN5SW8nRsiGKK7Aif4CrACPENkA6P/Lw6fHGQ==", + "dependencies": { + "@types/unist": "^3.0.0", + "unist-util-is": "^6.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/unist-util-is": { "version": "6.0.0", "resolved": "https://registry.npmjs.org/unist-util-is/-/unist-util-is-6.0.0.tgz", @@ -13739,6 +13998,19 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/unist-util-remove-position": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/unist-util-remove-position/-/unist-util-remove-position-5.0.0.tgz", + "integrity": "sha512-Hp5Kh3wLxv0PHj9m2yZhhLt58KzPtEYKQQ4yxfYFEO7EvHwzyDYnduhHnY1mDxoqr7VUwVuHXk9RXKIiYS1N8Q==", + "dependencies": { + "@types/unist": "^3.0.0", + "unist-util-visit": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/unist-util-stringify-position": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/unist-util-stringify-position/-/unist-util-stringify-position-4.0.0.tgz", @@ -13886,6 +14158,19 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/vfile-location": { + "version": "5.0.3", + "resolved": "https://registry.npmjs.org/vfile-location/-/vfile-location-5.0.3.tgz", + "integrity": "sha512-5yXvWDEgqeiYiBe1lbxYF7UMAIm/IcopxMHrMQDq3nvKcjPKIhZklUKL+AE7J7uApI4kwe2snsK+eI6UTj9EHg==", + "dependencies": { + "@types/unist": "^3.0.0", + "vfile": "^6.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/vfile-message": { "version": "4.0.3", "resolved": "https://registry.npmjs.org/vfile-message/-/vfile-message-4.0.3.tgz", @@ -14125,6 +14410,15 @@ "node": ">=18" } }, + "node_modules/web-namespaces": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/web-namespaces/-/web-namespaces-2.0.1.tgz", + "integrity": "sha512-bKr1DkiNa2krS7qxNtdrtHAmzuYGFQLiQ13TsorsdT6ULTkPLKuu5+GsFpDlg6JFjUTwX2DyhMPG2be8uPrqsQ==", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/wooorm" + } + }, "node_modules/webidl-conversions": { "version": "7.0.0", "resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-7.0.0.tgz", diff --git a/package.json b/package.json index 68de2aa..013f336 100644 --- a/package.json +++ b/package.json @@ -19,6 +19,7 @@ "fft.js": "^4.0.4", "idb": "^8.0.3", "jszip": "^3.10.1", + "katex": "^0.16.46", "onnxruntime-web": "^1.26.0", "openai": "^6.33.0", "react": "^19.1.0", @@ -27,7 +28,9 @@ "react-markdown": "^10.1.0", "react-syntax-highlighter": "^15.6.1", "reflect-metadata": "^0.2.2", + "rehype-katex": "^7.0.1", "remark-gfm": "^4.0.1", + "remark-math": "^6.0.0", "tone": "^15.1.22", "vexflow": "^5.0.0", "zustand": "^5.0.6" diff --git a/src/components/ChatBox.css b/src/components/ChatBox.css index 53e5f8c..f71ee3a 100644 --- a/src/components/ChatBox.css +++ b/src/components/ChatBox.css @@ -278,6 +278,16 @@ color: #ccc; } +.message-content .katex-display { + margin: 8px 0; + overflow-x: auto; + overflow-y: hidden; +} + +.message-content .katex { + max-width: 100%; +} + .message-content table { border-collapse: collapse; width: 100%; diff --git a/src/components/chat/AssistantMessage.test.tsx b/src/components/chat/AssistantMessage.test.tsx index 04d68b2..567eecd 100644 --- a/src/components/chat/AssistantMessage.test.tsx +++ b/src/components/chat/AssistantMessage.test.tsx @@ -1,8 +1,12 @@ -import { fireEvent, render, screen } from '@testing-library/react'; -import { describe, expect, it, vi } from 'vitest'; +import { act, fireEvent, render, screen } from '@testing-library/react'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import AssistantMessage from './AssistantMessage'; describe('AssistantMessage', () => { + afterEach(() => { + vi.useRealTimers(); + }); + it.each([ 'Thinking... click here to abort.', 'Processing... 3 tokens received. click here to abort.' @@ -22,4 +26,71 @@ describe('AssistantMessage', () => { fireEvent.click(abortButton); expect(onAbort).toHaveBeenCalledTimes(1); }); + + it('shows and updates the thinking timer while waiting for tokens', () => { + vi.useFakeTimers(); + + render( + Thinking... click here to abort.'} + isStreaming + onAbort={vi.fn()} + /> + ); + + expect(screen.getByText('Thinking for 0s...')).toBeInTheDocument(); + + act(() => { + vi.advanceTimersByTime(12_000); + }); + + expect(screen.getByText('Thinking for 12s...')).toBeInTheDocument(); + }); + + it('shows minute formatting after one minute of thinking', () => { + vi.useFakeTimers(); + + render( + Thinking... click here to abort.'} + isStreaming + onAbort={vi.fn()} + /> + ); + + act(() => { + vi.advanceTimersByTime(65_000); + }); + + expect(screen.getByText('Thinking for 1m 05s...')).toBeInTheDocument(); + }); + + it('renders inline LaTeX with KaTeX markup', () => { + const { container } = render( + + ); + + expect(container.querySelector('.katex')).toBeInTheDocument(); + expect(container.querySelector('.katex-mathml')).toBeInTheDocument(); + expect(screen.queryByText('$\\rightarrow$')).not.toBeInTheDocument(); + }); + + it('renders block LaTeX as display math', () => { + const { container } = render( + + ); + + expect(container.querySelector('.katex-display')).toBeInTheDocument(); + }); + + it('renders markdown code blocks alongside LaTeX', () => { + const { container } = render( + + ); + + expect(container.querySelector('.katex')).toBeInTheDocument(); + const codeElement = container.querySelector('code.language-ts'); + expect(codeElement).toBeInTheDocument(); + expect(codeElement).toHaveTextContent('const value = 1;'); + }); }); diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx index 96b2b93..4573f46 100644 --- a/src/components/chat/AssistantMessage.tsx +++ b/src/components/chat/AssistantMessage.tsx @@ -1,6 +1,8 @@ -import React, { memo } from 'react'; +import React, { memo, useEffect, useState } from 'react'; import ReactMarkdown from 'react-markdown'; +import rehypeKatex from 'rehype-katex'; import remarkGfm from 'remark-gfm'; +import remarkMath from 'remark-math'; import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism'; import type { PerformanceInfo } from '../../agent/llm/StreamingTypes'; @@ -40,11 +42,45 @@ const formatTps = (value?: number): string | null => { return value.toFixed(1); }; +const THINKING_LABEL = 'Thinking...'; +const PROCESSING_LABEL = 'Processing...'; + +const formatThinkingDuration = (elapsedSeconds: number): string => { + if (elapsedSeconds < 60) { + return `Thinking for ${elapsedSeconds}s...`; + } + + const minutes = Math.floor(elapsedSeconds / 60); + const seconds = elapsedSeconds % 60; + return `Thinking for ${minutes}m ${seconds.toString().padStart(2, '0')}s...`; +}; + const AssistantMessage: React.FC = ({ content, isStreaming, onAbort, performanceInfo }) => { const prefillTps = formatTps(performanceInfo?.prefillTps); const generationTps = formatTps(performanceInfo?.generationTps); const hasPerformanceInfo = Boolean(prefillTps || generationTps); - const processingWaveLabels = ['Thinking...', 'Processing...']; + const [thinkingElapsedSeconds, setThinkingElapsedSeconds] = useState(0); + const processingWaveLabels = [THINKING_LABEL, PROCESSING_LABEL]; + const isThinking = isStreaming && content.includes(`${THINKING_LABEL}`); + + useEffect(() => { + if (!isThinking) { + setThinkingElapsedSeconds(0); + return; + } + + setThinkingElapsedSeconds(0); + + const startedAt = Date.now(); + const intervalId = window.setInterval(() => { + const elapsedSeconds = Math.floor((Date.now() - startedAt) / 1000); + setThinkingElapsedSeconds(elapsedSeconds); + }, 1000); + + return () => { + window.clearInterval(intervalId); + }; + }, [isThinking]); const renderContent = () => { // Handle special abort link for streaming messages @@ -59,7 +95,9 @@ const AssistantMessage: React.FC = ({ content, isStreamin processingWaveMarkup, '' ); - const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...'; + const waveLabel = processingWaveMarkup.includes(THINKING_LABEL) + ? formatThinkingDuration(thinkingElapsedSeconds) + : processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? THINKING_LABEL; return ( @@ -87,7 +125,8 @@ const AssistantMessage: React.FC = ({ content, isStreamin return ( Date: Fri, 15 May 2026 12:30:12 -0700 Subject: [PATCH 18/19] fix: failed tests --- src/agent/core/SystemPrompts.ts | 2 +- src/util/opfsModelCache.ts | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/agent/core/SystemPrompts.ts b/src/agent/core/SystemPrompts.ts index 9f4e80d..ad292da 100644 --- a/src/agent/core/SystemPrompts.ts +++ b/src/agent/core/SystemPrompts.ts @@ -262,6 +262,6 @@ export class SystemPrompts { * Clear the cached template (useful for development/testing) */ static clearCache(): void { - this.cachedTemplate = null; + this.cachedTemplates.clear(); } } diff --git a/src/util/opfsModelCache.ts b/src/util/opfsModelCache.ts index 6d37198..1b6202e 100644 --- a/src/util/opfsModelCache.ts +++ b/src/util/opfsModelCache.ts @@ -123,13 +123,14 @@ export class OpfsModelCache { const finalWritable = await finalHandle.createWritable(); try { const tempFile = await tempHandle.getFile(); + const tempBuffer = await tempFile.arrayBuffer(); console.log('[opfsModelCache] Temp file ready for finalize copy.', { filename, tempFilename, tempSize: tempFile.size, expectedSize: sizeValue, }); - await finalWritable.write(tempFile); + await finalWritable.write(tempBuffer); await finalWritable.close(); } catch (error) { await finalWritable.abort(); From 1fd2c283f28bd7da3af1805a065e8f0524563df0 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Fri, 15 May 2026 12:34:18 -0700 Subject: [PATCH 19/19] docs: updated README.md and LICENSE --- LICENSE | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++-- README.md | 20 ++++++++++------ 2 files changed, 83 insertions(+), 9 deletions(-) diff --git a/LICENSE b/LICENSE index 6f20539..6bf1f7b 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ # License -Copyright 2025 Xiaohan Tian +Copyright 2025-2026 Xiaohan Tian Licensed under the Apache License, Version 2.0 (the "License"); you may not use this project except in compliance with the License. @@ -17,7 +17,7 @@ You may obtain a copy of the License at: If you use this project in any public or commercial product, you must provide the following credit in a prominent location (such as an "About" screen or documentation): > Powered by **K.G.Studio** - > © 2025 Xiaohan Tian + > © 2025-2026 Xiaohan Tian ## Standard Terms @@ -159,4 +159,72 @@ Modifications: Substantial textual and functional changes have been made; only t --- +### 5. Gemma 4 E4B (LiteRT-LM) + +This project uses the Gemma 4 E4B model by Google LLC, accessed via LiteRT-LM. +Gemma is licensed under the Apache License, Version 2.0, subject to the +Google Gemma Terms of Use. + +Original model: https://huggingface.co/google/gemma-4-E4B-it +Gemma Terms of Use: https://ai.google.dev/gemma/terms +Prohibited Use Policy: https://ai.google.dev/gemma/prohibited_use_policy + +Gemma is provided under and subject to the Gemma Terms of Use found at +ai.google.dev/gemma/terms. + +Copyright 2025 Google LLC + +--- + +### 6. UVR-MDX-NET-Inst_HQ_3 + +This project uses the UVR-MDX-NET-Inst_HQ_3 model from the Ultimate Vocal +Remover (UVR) project, distributed as an ONNX model via +https://huggingface.co/notabilia/uvr5-models. + +Ultimate Vocal Remover — https://github.com/Anjok07/ultimatevocalremovergui +Authors: Anjok07, KimberleyJSN, and the UVR development community + +MDX-NET architecture — https://github.com/kuielab/mdx-net +Authors: Woosung Choi, Minseok Kim, Jaehwa Chung, Daewon Lee, Soonyoung Jung (KUIELab) + +``` +MIT License + +Copyright (c) Anjok07, KimberleyJSN, and the UVR contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +``` + +--- + +### 7. MediaPipe / LiteRT (GenAI runtime) + +This project bundles the MediaPipe GenAI WebAssembly runtime +(`genai_bundle.mjs`, `genai_wasm_internal.js/wasm`, +`genai_wasm_module_internal.js/wasm`) from Google's MediaPipe project. + +MediaPipe is licensed under the Apache License, Version 2.0. + +Original project: https://github.com/google-ai-edge/mediapipe +Copyright 2019 The MediaPipe Authors + +--- + Apache License Version 2.0, January 2004: http://www.apache.org/licenses/LICENSE-2.0 \ No newline at end of file diff --git a/README.md b/README.md index be1435c..892d223 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,13 @@ # K.G.Studio — A Browser-based DAW with AI Assistant + + ## New! -> ### ✦ [**K.G.One Music Studio is available now.**](https://github.com/KGAudioLab/K.G.One) ✦
+> ### ✦ [K.G.One Music Studio is available now.](https://github.com/KGAudioLab/K.G.One) ✦
> [**K.G.One Music Studio**](https://github.com/KGAudioLab/K.G.One) is a fully local, open-source integrated platform built around **K.G.Studio** (this project). It bundles [ACE-Step 1.5](https://github.com/ace-step/ACE-Step-1.5) for full-song generation, [Foundation-1](https://huggingface.co/RoyalCities/Foundation-1) for audio clip and MIDI loop generation, and [python-audio-separator (UVR5)](https://github.com/nomadkaraoke/python-audio-separator) for stem separation — bringing GPU-accelerated AI music generation directly into your browser-based production workflow. ## What is K.G.Studio? @@ -19,10 +23,12 @@ K.G.Studio is a lightweight, modern DAW that runs entirely in the browser with * K.G.One Logo
-> Note: Full-Song Generation Feature, Audio Clip Generation Feature, and Stem Separation Feature requires [**K.G.One Music Studio**](https://github.com/KGAudioLab/K.G.One) integration. +> Note: Full-Song Generation Feature and Audio Clip Generation Feature requires [**K.G.One Music Studio**](https://github.com/KGAudioLab/K.G.One) integration. ## Latest Updates +- **2026.05.15**: Added **browser-embedded AI models** — two AI models now run entirely in the browser with no external service, no API key, and no K.G.One server required. The **K.G.Studio Musician Assistant** gains a new **Local LLM (Browser)** provider powered by **Gemma 4 E4B** via LiteRT-LM with WebGPU acceleration; the model is downloaded once and cached in OPFS for instant subsequent launches, with configurable context length (32 k / 64 k / 128 k tokens) and live inference performance statistics. **Stem separation** now also runs locally through a browser-embedded **UVR-MDX-NET-Inst_HQ_3** ONNX model with WebGPU acceleration — open the **Music Generator** panel (✦ button), download the model once, and separate vocals from instruments entirely on-device. Both features require a WebGPU-capable browser (Chrome 113+ or Edge 113+) and a secure context (HTTPS or localhost). Recommended hardware: a GPU with at least 8 GB VRAM or a system with at least 16 GB unified RAM. + - **2026.05.10**: Added **staff notation (sheet music) view** — the piano roll now offers a full standard notation mode. Switch between Piano Roll and Sheet Music views using the toggle in the piano roll toolbar. In sheet music mode, notes are engraved via VexFlow with automatic clef selection (treble or bass) based on the active instrument, key signature rendering, automatic beam grouping, ties across bar lines, and configurable quantization for note-value resolution. Enable **Track Scope** to render all MIDI regions on the track as a continuous score rather than a single isolated region. - **2026.05.09**: Added **audio recording** — record directly from your microphone into an audio track. A live waveform preview grows in real time as you record, and the region is committed to the timeline as a standard audio region when you stop. Added **audio I/O device selection** in Settings so you can choose your preferred microphone input and audio output device. @@ -47,8 +53,6 @@ K.G.Studio is a lightweight, modern DAW that runs entirely in the browser with * This project investigates how AI-human collaboration can enhance creative music-making, from intelligent harmony suggestions to automated editing tasks. As an experimental platform, expect frequent changes, evolving features, and occasional instability as we push the boundaries of what's possible in AI-assisted music production. -### Start using the app online: [K.G.Studio (kgaudiolab.github.io/kgstudio)](https://kgaudiolab.github.io/kgstudio) - ## Demo Videos
@@ -329,15 +333,17 @@ Feature priorities might change. - [X] Recording - [X] Event List - [X] Add support for OpenAI's open source models (`gpt-oss-20b` and `gpt-oss-120b`) +- [X] Stuff notation +- [X] K.G.One Music Studio integration +- [X] Browser-embedded AI models (on-device LLM via Gemma 4 E4B; on-device stem separation via UVR-MDX-NET-Inst_HQ_3) ### Post 1.0 -- [ ] Stuff notation +- [ ] EQ - [ ] Filters and effects - [ ] Virtual MIDI device output - [ ] Enhanced AI Music Assistant Agent - ## Help Needed We're looking for contributors to help make K.G.Studio even better! Whether you're a developer, musician, or designer, your expertise can make a real difference. @@ -385,4 +391,4 @@ Licensed under the Apache License, Version 2.0, with additional terms (see `LICE - No patent applications using this software or assets - Attribution required when used in public/commercial products (“Powered by K.G.Studio”) -Third‑party notices (e.g., FluidR3_GM SoundFont, midi‑js‑soundfonts, VexFlow, and prompt structure notes) are included in `LICENSE`. +Third‑party notices (e.g., FluidR3_GM SoundFont, midi‑js‑soundfonts, VexFlow, prompt structure notes, Gemma 4 E4B, UVR-MDX-NET-Inst_HQ_3, and MediaPipe) are included in `LICENSE`.