3dtours/KGStudio
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: added local separator model demucs_4s support

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
src/components/KGOnePanel.test.tsx
+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();
});
});
});