feat: added local separator model demucs_4s support

This commit is contained in:
Xiaohan-Tian
2026-05-27 20:17:17 -07:00
parent e2a7d1c285
commit 2140e944f2
18 changed files with 723 additions and 149 deletions
+51 -16
View File
@@ -4,6 +4,7 @@ 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';
import { LOCAL_SEPARATOR_MODEL_IDS } from '../util/local-separator/config';
const { mockLocalSeparatorDownload } = vi.hoisted(() => ({
mockLocalSeparatorDownload: vi.fn(async (_url?: string, _filename?: string, _onProgress?: unknown) => undefined),
@@ -11,7 +12,7 @@ const { mockLocalSeparatorDownload } = vi.hoisted(() => ({
let kgoneEnabled = false;
let selectedRegionIds: string[] = [];
let localModelCached = false;
let localModelCached: Record<string, boolean> = {};
let localSeparationResult: Array<{ name: string; blob: Blob }> = [];
const mockRefreshProjectState = vi.fn();
@@ -36,6 +37,7 @@ vi.mock('../core/config/ConfigManager', () => ({
if (key === 'general.kgone.enabled') return kgoneEnabled;
if (key === 'general.kgone.base_url') return 'http://127.0.0.1:8000';
if (key === 'general.uvr5_web_runtime.mdx_net_model_url') return 'https://example.com/custom-uvr5.onnx';
if (key === 'general.uvr5_web_runtime.htdemucs_4s_model_url') return 'https://example.com/custom-htdemucs.onnx';
return undefined;
},
}),
@@ -83,13 +85,13 @@ vi.mock('../util/audioUtil', () => ({
vi.mock('../util/local-separator/modelCache', () => ({
LocalSeparatorModelCache: {
exists: vi.fn(async () => localModelCached),
download: vi.fn(async (url: string, filename: string, onProgress: (progress: unknown) => void) => {
localModelCached = true;
return mockLocalSeparatorDownload(url, filename, onProgress);
exists: vi.fn(async (modelConfig: { id: string }) => localModelCached[modelConfig.id] ?? false),
download: vi.fn(async (modelConfig: { id: string; filename: string }, url: string, onProgress: (progress: unknown) => void) => {
localModelCached[modelConfig.id] = true;
return mockLocalSeparatorDownload(modelConfig.filename, url, onProgress);
}),
delete: vi.fn(async () => {
localModelCached = false;
delete: vi.fn(async (modelConfig: { id: string }) => {
localModelCached[modelConfig.id] = false;
}),
getArrayBuffer: vi.fn(async () => new ArrayBuffer(16)),
},
@@ -125,7 +127,7 @@ describe('KGOnePanel local separator mode', () => {
beforeEach(() => {
kgoneEnabled = false;
selectedRegionIds = [];
localModelCached = false;
localModelCached = {};
localSeparationResult = [
{ name: 'Instrumental', blob: new Blob(['instrumental'], { type: 'audio/wav' }) },
{ name: 'Vocals', blob: new Blob(['vocals'], { type: 'audio/wav' }) },
@@ -143,41 +145,42 @@ describe('KGOnePanel local separator mode', () => {
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();
expect(screen.getByRole('button', { name: 'Download Selected Model' })).toBeInTheDocument();
});
it('shows the single local separator model and advanced settings when the model is cached', async () => {
localModelCached = true;
localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
selectedRegionIds = ['audio-region-1'];
render(<KGOnePanel isVisible={true} />);
await screen.findByText('Selected Region');
const options = await screen.findAllByRole('option');
expect(options).toHaveLength(1);
expect(options).toHaveLength(2);
expect(options[0]).toHaveTextContent('Vocal and Instrument (Medium Accuracy)');
expect(options[1]).toHaveTextContent('Vocal, Drums, Bass, and Others');
fireEvent.click(screen.getByRole('button', { name: /Advanced Settings/i }));
expect(screen.getByLabelText('Optional audio chunk duration (seconds)')).toBeInTheDocument();
expect(screen.getByLabelText('MDX overlap')).toBeInTheDocument();
expect(screen.getByLabelText('Model overlap')).toBeInTheDocument();
});
it('uses the configured UVR5 model URL when downloading the local model', async () => {
render(<KGOnePanel isVisible={true} />);
fireEvent.click(await screen.findByRole('button', { name: 'Download Model' }));
fireEvent.click(await screen.findByRole('button', { name: 'Download Selected Model' }));
await waitFor(() => {
expect(mockLocalSeparatorDownload).toHaveBeenCalledWith(
'https://example.com/custom-uvr5.onnx',
'UVR-MDX-NET-Inst_HQ_3.onnx',
'https://example.com/custom-uvr5.onnx',
expect.any(Function),
);
});
});
it('prompts for an audio region when the model is cached but nothing is selected', async () => {
localModelCached = true;
localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
render(<KGOnePanel isVisible={true} />);
@@ -185,7 +188,7 @@ describe('KGOnePanel local separator mode', () => {
});
it('renders local separation outputs after processing completes', async () => {
localModelCached = true;
localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
selectedRegionIds = ['audio-region-1'];
render(<KGOnePanel isVisible={true} />);
@@ -198,4 +201,36 @@ describe('KGOnePanel local separator mode', () => {
expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument();
});
});
it('uses Demucs defaults and renders four local stem players', async () => {
localModelCached[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s] = true;
localSeparationResult = [
{ name: 'Vocals', blob: new Blob(['vocals'], { type: 'audio/wav' }) },
{ name: 'Drums', blob: new Blob(['drums'], { type: 'audio/wav' }) },
{ name: 'Bass', blob: new Blob(['bass'], { type: 'audio/wav' }) },
{ name: 'Others', blob: new Blob(['others'], { type: 'audio/wav' }) },
];
selectedRegionIds = ['audio-region-1'];
render(<KGOnePanel isVisible={true} />);
await screen.findByText('Selected Region');
fireEvent.change(screen.getByRole('combobox'), { target: { value: LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s } });
fireEvent.click(screen.getByRole('button', { name: /Advanced Settings/i }));
await waitFor(() => {
expect((screen.getByLabelText('Optional audio chunk duration (seconds)') as HTMLInputElement).value).toBe('8');
expect((screen.getByLabelText('Model overlap') as HTMLInputElement).value).toBe('0.25');
});
fireEvent.click(screen.getByRole('button', { name: 'Separate Stems' }));
await waitFor(() => {
expect(screen.getByText('Vocals')).toBeInTheDocument();
expect(screen.getByText('Drums')).toBeInTheDocument();
expect(screen.getByText('Bass')).toBeInTheDocument();
expect(screen.getByText('Others')).toBeInTheDocument();
expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument();
});
});
});
+57 -35
View File
@@ -16,13 +16,15 @@ 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_DEFAULT_MODEL_URL,
getLocalSeparatorModelConfig,
LOCAL_SEPARATOR_MODELS,
LOCAL_SEPARATOR_MODEL_CONFIGS,
LOCAL_SEPARATOR_MODEL_IDS,
} from '../util/local-separator/config';
import { LocalSeparatorModelCache } from '../util/local-separator/modelCache';
import { runLocalSeparator } from '../util/local-separator/runner';
import { LocalOrtRuntimeManager, detectLocalRuntimeSupport } from '../util/local-separator/runtime';
import type { LocalSeparatorModelConfig, LocalSeparatorModelId } from '../util/local-separator/types';
// ─── Types ────────────────────────────────────────────────────────────────────
@@ -36,9 +38,10 @@ const SERVER_SEPARATOR_MODELS = [
{ 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 LOCAL_SEPARATOR_MODEL_OPTIONS = LOCAL_SEPARATOR_MODELS.map(modelConfig => ({
label: modelConfig.displayName,
value: modelConfig.id,
})) as ReadonlyArray<{ label: string; value: LocalSeparatorModelId }>;
const KGONE_TABS = ['fullsong', 'remix', 'repaint', 'separator'] as const;
const CLIP_NOTES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'];
@@ -862,8 +865,11 @@ function countRepaintTracks(sourceTrackName: string): number {
const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
const { selectedRegionIds, projectName, bpm, timeSignature, maxBars, refreshProjectState } = useProjectStore();
const localOnlyMode = mode === 'local-separator';
const availableSeparatorModels = localOnlyMode ? LOCAL_SEPARATOR_MODELS : SERVER_SEPARATOR_MODELS;
const [model, setModel] = useState<typeof SERVER_SEPARATOR_MODELS[number]['value']>(availableSeparatorModels[0].value);
const availableSeparatorModels = localOnlyMode ? LOCAL_SEPARATOR_MODEL_OPTIONS : SERVER_SEPARATOR_MODELS;
const [model, setModel] = useState<string>(availableSeparatorModels[0].value);
const currentLocalModelConfig = useMemo<LocalSeparatorModelConfig>(() => {
return getLocalSeparatorModelConfig(model);
}, [model]);
// Generation state
const [genStatus, setGenStatus] = useState<GenStatus>('idle');
@@ -881,7 +887,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
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 [localOverlap, setLocalOverlap] = useState(String(currentLocalModelConfig.defaults.overlap));
const abortRef = useRef<AbortController | null>(null);
const taskIdRef = useRef<string>('');
@@ -922,18 +928,28 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
setModel(availableSeparatorModels[0].value);
}, [availableSeparatorModels]);
useEffect(() => {
if (!localOnlyMode) return;
setLocalChunkDurationSeconds(
currentLocalModelConfig.defaultChunkDurationSeconds == null
? ''
: String(currentLocalModelConfig.defaultChunkDurationSeconds),
);
setLocalOverlap(String(currentLocalModelConfig.defaults.overlap));
}, [currentLocalModelConfig, localOnlyMode]);
const refreshLocalModelCacheState = useCallback(async () => {
if (!localOnlyMode) return;
setIsCheckingLocalModel(true);
try {
setIsLocalModelCached(await LocalSeparatorModelCache.exists());
setIsLocalModelCached(await LocalSeparatorModelCache.exists(currentLocalModelConfig));
} catch (err) {
console.error('[KGOne] Local model cache check failed:', err);
setErrorMsg(err instanceof Error ? err.message : String(err));
} finally {
setIsCheckingLocalModel(false);
}
}, [localOnlyMode]);
}, [currentLocalModelConfig, localOnlyMode]);
useEffect(() => {
if (!localOnlyMode) return;
@@ -960,11 +976,11 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
}, [selectedRegionIds]);
const getConfiguredLocalSeparatorModelUrl = useCallback(() => {
const configured = ConfigManager.instance().get('general.uvr5_web_runtime.mdx_net_model_url');
const configured = ConfigManager.instance().get(currentLocalModelConfig.download.configKey);
return typeof configured === 'string' && configured.trim()
? configured
: LOCAL_SEPARATOR_DEFAULT_MODEL_URL;
}, []);
: currentLocalModelConfig.download.defaultUrl;
}, [currentLocalModelConfig]);
const isGenerating = genStatus !== 'idle' && genStatus !== 'done' && genStatus !== 'error';
@@ -972,25 +988,25 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
setIsDownloadingLocalModel(true);
setErrorMsg('');
setLocalProgressPercent(0);
setLocalProgressText('Downloading local separator model...');
setLocalProgressText(`Downloading ${currentLocalModelConfig.displayName}...`);
try {
await LocalSeparatorModelCache.download(
currentLocalModelConfig,
getConfiguredLocalSeparatorModelUrl(),
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`,
? `Downloading ${currentLocalModelConfig.displayName}... ${receivedMb} / ${totalMb} MB`
: `Downloading ${currentLocalModelConfig.displayName}... ${receivedMb} MB`,
);
},
);
setLocalProgressPercent(100);
setLocalProgressText('Local separator model is ready.');
setLocalProgressText(`${currentLocalModelConfig.displayName} is ready.`);
await refreshLocalModelCacheState();
} catch (err) {
setLocalProgressPercent(0);
@@ -999,13 +1015,13 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
} finally {
setIsDownloadingLocalModel(false);
}
}, [getConfiguredLocalSeparatorModelUrl, refreshLocalModelCacheState]);
}, [currentLocalModelConfig, getConfiguredLocalSeparatorModelUrl, refreshLocalModelCacheState]);
const handleDeleteLocalModel = useCallback(async () => {
setIsDeletingLocalModel(true);
setErrorMsg('');
try {
await LocalSeparatorModelCache.delete();
await LocalSeparatorModelCache.delete(currentLocalModelConfig);
localRuntimeManagerRef.current?.reset();
setLocalProviderLabel(runtimeSupport.webgpuExposed ? 'webgpu available' : 'cpu/wasm only');
setLocalProgressPercent(0);
@@ -1016,7 +1032,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
} finally {
setIsDeletingLocalModel(false);
}
}, [refreshLocalModelCacheState, runtimeSupport.webgpuExposed]);
}, [currentLocalModelConfig, refreshLocalModelCacheState, runtimeSupport.webgpuExposed]);
const handleSeparateServer = useCallback(async () => {
if (!selectedAudioRegion) return;
@@ -1212,12 +1228,12 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
setLocalProviderLabel(runtimeSupport.webgpuExposed ? 'webgpu available' : 'cpu/wasm only');
try {
const modelBuffer = await LocalSeparatorModelCache.getArrayBuffer();
const modelBuffer = await LocalSeparatorModelCache.getArrayBuffer(currentLocalModelConfig);
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));
const runtime = await runtimeManager.ensureRuntime(currentLocalModelConfig, new Uint8Array(modelBuffer));
setGenStatus('generating');
setLocalProgressPercent(3);
@@ -1247,10 +1263,10 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
const result = await runLocalSeparator({
session: runtime.session,
runtimeProvider: runtime.provider,
modelConfig: LOCAL_SEPARATOR_MODEL_CONFIG,
modelConfig: currentLocalModelConfig,
audioBuffer: inputBuffer,
chunkDurationSeconds: Number.isFinite(chunkDuration) && (chunkDuration ?? 0) > 0 ? chunkDuration : null,
overlap: Number.isFinite(overlapValue) ? overlapValue : LOCAL_SEPARATOR_MODEL_CONFIG.defaults.overlap,
overlap: Number.isFinite(overlapValue) ? overlapValue : currentLocalModelConfig.defaults.overlap,
onProviderChange: provider => setLocalProviderLabel(provider),
onProgress: progress => {
setLocalProgressPercent(progress.percent);
@@ -1283,6 +1299,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
projectName,
localChunkDurationSeconds,
localOverlap,
currentLocalModelConfig,
]);
const handleSeparate = useCallback(async () => {
@@ -1381,14 +1398,19 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
<div className="kgone-local-mode-card">
<div className="kgone-local-mode-title">Local Separator Mode</div>
<div className="kgone-local-mode-text">
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.{' '}
If K.G.One Music Studio is unavailable, Local Separator Mode provides a built-in alternative for extracting stems
directly in your browser. Two local models are available: Vocal and Instrument (Medium Accuracy), and Vocal,
Drums, Bass, and Others. Download status below reflects the currently selected model. Vocal and Instrument
(Medium Accuracy) usually takes longer to process than Vocal, Drums, Bass, and Others, and total processing time
will still depend on your hardware. If processing falls back to CPU, the page may become temporarily less
responsive while separation is running.{' '}
<a href="https://github.com/KGAudioLab/K.G.One" target="_blank" rel="noopener noreferrer">Learn more about K.G.One Music Studio server integration.</a>
</div>
<div className="kgone-runtime-row">
<div className="kgone-provider-chip">Provider: {localProviderLabel}</div>
<div className="kgone-provider-chip">Model: {isLocalModelCached ? 'downloaded' : 'not downloaded'}</div>
<div className="kgone-provider-chip">
Model: {currentLocalModelConfig.displayName} ({isLocalModelCached ? 'downloaded' : 'not downloaded'})
</div>
</div>
{(localProgressText || isCheckingLocalModel) && (
<div className="kgone-progress-block">
@@ -1414,7 +1436,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
disabled={isCheckingLocalModel || isDownloadingLocalModel || isDeletingLocalModel || isGenerating}
onClick={() => void handleDownloadLocalModel()}
>
{isDownloadingLocalModel ? 'Downloading Model...' : 'Download Model'}
{isDownloadingLocalModel ? 'Downloading Model...' : 'Download Selected Model'}
</button>
) : (
<>
@@ -1451,7 +1473,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
<div className="kgone-field">
<label className="kgone-label">Separation Model</label>
<select className="kgone-select" value={model} onChange={e => setModel(e.target.value as typeof SERVER_SEPARATOR_MODELS[number]['value'])}>
<select className="kgone-select" value={model} onChange={e => setModel(e.target.value)}>
{availableSeparatorModels.map(m => (
<option key={m.value} value={m.value}>{m.label}</option>
))}
@@ -1474,10 +1496,10 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
/>
</div>
<div className="kgone-field">
<label className="kgone-label">MDX overlap</label>
<label className="kgone-label">Model overlap</label>
<input
className="kgone-input"
aria-label="MDX overlap"
aria-label="Model overlap"
type="number"
min={0.001}
max={0.999}
@@ -1552,7 +1574,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
) : (
<div className="kgone-separator-hint">
{localOnlyMode && !isLocalModelCached
? 'Download the local separator model, then select an audio region on the timeline to extract stems from it.'
? `Download ${currentLocalModelConfig.displayName}, 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.'}
</div>
)}
@@ -29,6 +29,7 @@ const configState = new Map<string, unknown>([
['general.local_browser.context_length', 65536],
['general.local_browser.model_url', 'https://huggingface.co/notabilia/gemma-4-E4B-it-litert-lm/resolve/main/gemma-4-E4B-it-web.task'],
['general.uvr5_web_runtime.mdx_net_model_url', 'https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx'],
['general.uvr5_web_runtime.htdemucs_4s_model_url', 'https://huggingface.co/notabilia/uvr5-models/resolve/main/htdemucs_embedded.onnx'],
['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'],
@@ -140,6 +141,7 @@ describe('GeneralSettings', () => {
expect(await screen.findByDisplayValue('https://huggingface.co/notabilia/gemma-4-E4B-it-litert-lm/resolve/main/gemma-4-E4B-it-web.task')).toBeTruthy();
expect(screen.getByDisplayValue('https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx')).toBeTruthy();
expect(screen.getByDisplayValue('https://huggingface.co/notabilia/uvr5-models/resolve/main/htdemucs_embedded.onnx')).toBeTruthy();
const inputs = screen.getAllByRole('textbox');
const gemmaUrlInput = inputs.find(input =>
@@ -148,22 +150,30 @@ describe('GeneralSettings', () => {
const uvr5UrlInput = inputs.find(input =>
(input as HTMLInputElement).value.includes('UVR-MDX-NET-Inst_HQ_3.onnx'),
) as HTMLInputElement | undefined;
const htdemucsUrlInput = inputs.find(input =>
(input as HTMLInputElement).value.includes('htdemucs_embedded.onnx'),
) as HTMLInputElement | undefined;
expect(gemmaUrlInput).toBeTruthy();
expect(uvr5UrlInput).toBeTruthy();
expect(htdemucsUrlInput).toBeTruthy();
fireEvent.change(gemmaUrlInput!, { target: { value: 'https://example.com/gemma.task' } });
fireEvent.change(uvr5UrlInput!, { target: { value: 'https://example.com/uvr5.onnx' } });
fireEvent.change(htdemucsUrlInput!, { target: { value: 'https://example.com/htdemucs.onnx' } });
await waitFor(() => {
expect(configManagerMock.set).toHaveBeenCalledWith('general.local_browser.model_url', 'https://example.com/gemma.task');
expect(configManagerMock.set).toHaveBeenCalledWith('general.uvr5_web_runtime.mdx_net_model_url', 'https://example.com/uvr5.onnx');
expect(configManagerMock.set).toHaveBeenCalledWith('general.uvr5_web_runtime.htdemucs_4s_model_url', 'https://example.com/htdemucs.onnx');
});
});
it('restores default download URLs and deletes the UVR5 model cache', async () => {
localSeparatorModelCacheMock.exists
.mockResolvedValueOnce(true)
.mockResolvedValueOnce(true)
.mockResolvedValueOnce(false)
.mockResolvedValueOnce(false);
render(<GeneralSettings />);
@@ -173,6 +183,7 @@ describe('GeneralSettings', () => {
const restoreLinks = screen.getAllByText('Restore default');
fireEvent.click(restoreLinks[0]);
fireEvent.click(restoreLinks[1]);
fireEvent.click(restoreLinks[2]);
const uvr5DeleteButton = screen.getAllByRole('button', { name: 'Delete Cached Model' })[1];
expect(uvr5DeleteButton).not.toBeDisabled();
fireEvent.click(uvr5DeleteButton);
@@ -186,6 +197,10 @@ describe('GeneralSettings', () => {
'general.uvr5_web_runtime.mdx_net_model_url',
'https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx',
);
expect(configManagerMock.set).toHaveBeenCalledWith(
'general.uvr5_web_runtime.htdemucs_4s_model_url',
'https://huggingface.co/notabilia/uvr5-models/resolve/main/htdemucs_embedded.onnx',
);
expect(localSeparatorModelCacheMock.delete).toHaveBeenCalled();
});
@@ -12,7 +12,10 @@ import {
normalizeLocalLLMContextLength,
type LocalLLMContextLength,
} from '../../../util/localLLMConfig';
import { LOCAL_SEPARATOR_DEFAULT_MODEL_URL } from '../../../util/local-separator/config';
import {
LOCAL_SEPARATOR_MODEL_CONFIGS,
LOCAL_SEPARATOR_MODEL_IDS,
} from '../../../util/local-separator/config';
const GeneralSettings: React.FC = () => {
const [llmProvider, setLlmProvider] = useState<string>(LOCAL_LLM_PROVIDER_KEY);
@@ -39,9 +42,12 @@ const GeneralSettings: React.FC = () => {
const [localModelState, setLocalModelState] = useState<LocalLLMModelState>(LocalLLMModelManager.getState());
const [localModelUrl, setLocalModelUrl] = useState<string>('');
const [uvr5ModelUrl, setUvr5ModelUrl] = useState<string>('');
const [htdemucsModelUrl, setHtdemucsModelUrl] = useState<string>('');
const [isUvr5ModelCached, setIsUvr5ModelCached] = useState<boolean>(false);
const [isCheckingUvr5ModelCache, setIsCheckingUvr5ModelCache] = useState<boolean>(false);
const [isDeletingUvr5Model, setIsDeletingUvr5Model] = useState<boolean>(false);
const [isHtdemucsModelCached, setIsHtdemucsModelCached] = useState<boolean>(false);
const [isDeletingHtdemucsModel, setIsDeletingHtdemucsModel] = useState<boolean>(false);
const configManager = ConfigManager.instance();
@@ -62,10 +68,16 @@ const GeneralSettings: React.FC = () => {
const refreshUvr5ModelCacheState = useCallback(async () => {
setIsCheckingUvr5ModelCache(true);
try {
setIsUvr5ModelCached(await LocalSeparatorModelCache.exists());
const [mdxCached, demucsCached] = await Promise.all([
LocalSeparatorModelCache.exists(LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium]),
LocalSeparatorModelCache.exists(LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s]),
]);
setIsUvr5ModelCached(mdxCached);
setIsHtdemucsModelCached(demucsCached);
} catch (error) {
console.error('Failed to check UVR5 cached model state:', error);
setIsUvr5ModelCached(false);
setIsHtdemucsModelCached(false);
} finally {
setIsCheckingUvr5ModelCache(false);
}
@@ -95,7 +107,14 @@ const GeneralSettings: React.FC = () => {
setCompatibleModel((configManager.get('general.openai_compatible.model') as string) || '');
setLocalContextLength(normalizeLocalLLMContextLength(configManager.get('general.local_browser.context_length')));
setLocalModelUrl((configManager.get('general.local_browser.model_url') as string) || LOCAL_LLM_DEFAULT_MODEL_URL);
setUvr5ModelUrl((configManager.get('general.uvr5_web_runtime.mdx_net_model_url') as string) || LOCAL_SEPARATOR_DEFAULT_MODEL_URL);
setUvr5ModelUrl(
(configManager.get('general.uvr5_web_runtime.mdx_net_model_url') as string)
|| LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium].download.defaultUrl,
);
setHtdemucsModelUrl(
(configManager.get('general.uvr5_web_runtime.htdemucs_4s_model_url') as string)
|| LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s].download.defaultUrl,
);
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) || '');
@@ -245,6 +264,11 @@ const GeneralSettings: React.FC = () => {
debouncedSave('general.uvr5_web_runtime.mdx_net_model_url', value);
};
const handleHtdemucsModelUrlChange = (value: string) => {
setHtdemucsModelUrl(value);
debouncedSave('general.uvr5_web_runtime.htdemucs_4s_model_url', value);
};
const handleDeleteLocalModel = async () => {
try {
await LocalLLMModelManager.deleteCachedModel();
@@ -256,7 +280,7 @@ const GeneralSettings: React.FC = () => {
const handleDeleteUvr5Model = async () => {
setIsDeletingUvr5Model(true);
try {
await LocalSeparatorModelCache.delete();
await LocalSeparatorModelCache.delete(LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium]);
setIsUvr5ModelCached(false);
} catch (error) {
console.error('Failed to delete UVR5 cached model:', error);
@@ -266,6 +290,19 @@ const GeneralSettings: React.FC = () => {
}
};
const handleDeleteHtdemucsModel = async () => {
setIsDeletingHtdemucsModel(true);
try {
await LocalSeparatorModelCache.delete(LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s]);
setIsHtdemucsModelCached(false);
} catch (error) {
console.error('Failed to delete HTDemucs cached model:', error);
} finally {
setIsDeletingHtdemucsModel(false);
await refreshUvr5ModelCacheState();
}
};
const handleLocalContextLengthChange = async (value: string) => {
const parsed = Number(value);
const normalized = normalizeLocalLLMContextLength(parsed);
@@ -448,7 +485,7 @@ const GeneralSettings: React.FC = () => {
<input
type="text"
className="settings-input"
placeholder={`e.g. ${LOCAL_SEPARATOR_DEFAULT_MODEL_URL}`}
placeholder={`e.g. ${LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium].download.defaultUrl}`}
value={uvr5ModelUrl}
onChange={(e) => handleUvr5ModelUrlChange(e.target.value)}
/>
@@ -458,7 +495,9 @@ const GeneralSettings: React.FC = () => {
href="#"
onClick={(e) => {
e.preventDefault();
handleUvr5ModelUrlChange(LOCAL_SEPARATOR_DEFAULT_MODEL_URL);
handleUvr5ModelUrlChange(
LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium].download.defaultUrl,
);
}}
style={{ color: '#5a9fd4', textDecoration: 'underline', cursor: 'pointer' }}
>
@@ -477,6 +516,45 @@ const GeneralSettings: React.FC = () => {
{isDeletingUvr5Model ? 'Deleting...' : 'Delete Cached Model'}
</button>
</div>
<div className="settings-item">
<label className="settings-label">
htdemucs_4s Download URL
</label>
<input
type="text"
className="settings-input"
placeholder={`e.g. ${LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s].download.defaultUrl}`}
value={htdemucsModelUrl}
onChange={(e) => handleHtdemucsModelUrlChange(e.target.value)}
/>
<div className="settings-help" style={{ fontSize: '12px', color: '#888', marginTop: '4px' }}>
Changing this URL may break downloads or point to an incompatible model file.{' '}
<a
href="#"
onClick={(e) => {
e.preventDefault();
handleHtdemucsModelUrlChange(
LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s].download.defaultUrl,
);
}}
style={{ color: '#5a9fd4', textDecoration: 'underline', cursor: 'pointer' }}
>
Restore default
</a>
</div>
</div>
<div className="settings-item" style={{ marginTop: '12px' }}>
<button
type="button"
className="settings-btn settings-btn-danger"
onClick={() => void handleDeleteHtdemucsModel()}
disabled={isCheckingUvr5ModelCache || isDeletingHtdemucsModel || !isHtdemucsModelCached}
>
{isDeletingHtdemucsModel ? 'Deleting...' : 'Delete Cached Model'}
</button>
</div>
</div>
<div className="settings-group">