feat: added local separator model demucs_4s support

2026-05-27 20:17:17 -07:00
parent e2a7d1c285
commit 2140e944f2
18 changed files with 723 additions and 149 deletions
@@ -1,32 +1,81 @@
-import type { LocalSeparatorModelConfig } from './types';
+import type { LocalSeparatorModelConfig, LocalSeparatorModelId } from './types';

-export const LOCAL_SEPARATOR_DEFAULT_MODEL_URL =
-  'https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx';
+export const LOCAL_SEPARATOR_MODEL_IDS = {
+  mdxMedium: 'UVR-MDX-NET-Inst_HQ_3.onnx',
+  htdemucs4s: 'htdemucs_4s.onnx',
+} as const satisfies Record<string, LocalSeparatorModelId>;

-export const LOCAL_SEPARATOR_MODEL_FILENAME = 'UVR-MDX-NET-Inst_HQ_3.onnx';
-export const LOCAL_SEPARATOR_MODEL_EXPECTED_SIZE_BYTES = 66759214;
+export const LOCAL_SEPARATOR_DEFAULT_MODEL_ID = LOCAL_SEPARATOR_MODEL_IDS.mdxMedium;

-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,
+export const LOCAL_SEPARATOR_MODEL_CONFIGS: Record<LocalSeparatorModelId, LocalSeparatorModelConfig> = {
+  [LOCAL_SEPARATOR_MODEL_IDS.mdxMedium]: {
+    id: LOCAL_SEPARATOR_MODEL_IDS.mdxMedium,
+    filename: 'UVR-MDX-NET-Inst_HQ_3.onnx',
+    kind: 'mdx',
+    displayName: 'Vocal and Instrument (Medium Accuracy)',
+    status: 'ready',
+    outputStemNames: ['Instrumental', 'Vocals'],
+    defaultChunkDurationSeconds: null,
+    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',
+    },
+    download: {
+      configKey: 'general.uvr5_web_runtime.mdx_net_model_url',
+      defaultUrl: 'https://huggingface.co/notabilia/uvr5-models/resolve/main/UVR-MDX-NET-Inst_HQ_3.onnx',
+      expectedSizeBytes: 66759214,
+    },
  },
-  metadata: {
-    compensate: 1.021,
-    mdx_dim_f_set: 3072,
-    mdx_dim_t_set: 8,
-    mdx_n_fft_scale_set: 7680,
-    primary_stem: 'Instrumental',
+  [LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s]: {
+    id: LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s,
+    filename: 'htdemucs_4s.onnx',
+    kind: 'demucs',
+    displayName: 'Vocal, Drums, Bass, and Others',
+    status: 'ready',
+    outputStemNames: ['Vocals', 'Drums', 'Bass', 'Others'],
+    defaultChunkDurationSeconds: 8,
+    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: null,
+    download: {
+      configKey: 'general.uvr5_web_runtime.htdemucs_4s_model_url',
+      defaultUrl: 'https://huggingface.co/notabilia/uvr5-models/resolve/main/htdemucs_embedded.onnx',
+      expectedSizeBytes: 180534758,
+    },
  },
 };
+
+export const LOCAL_SEPARATOR_MODELS = Object.values(LOCAL_SEPARATOR_MODEL_CONFIGS);
+
+export function getLocalSeparatorModelConfig(modelId: string): LocalSeparatorModelConfig {
+  const config = LOCAL_SEPARATOR_MODEL_CONFIGS[modelId as LocalSeparatorModelId];
+  if (!config) {
+    return LOCAL_SEPARATOR_MODEL_CONFIGS[LOCAL_SEPARATOR_DEFAULT_MODEL_ID];
+  }
+  return config;
+}
@@ -1,4 +1,4 @@
-import type { LocalSeparatorModelConfig, StereoChannels } from './types';
+import type { LocalSeparatorMdxModelConfig, StereoChannels } from './types';
 import { FFT, createWindowCache, getHannPeriodic, index4d, reflectPad } from './shared';

 interface SpectrogramPayload {
@@ -18,7 +18,7 @@ export class LocalSeparatorCpuDsp {
  private readonly numFreqBins: number;
  private readonly inverseFft: FFT;

-  constructor(config: LocalSeparatorModelConfig) {
+  constructor(config: LocalSeparatorMdxModelConfig) {
    this.nFft = config.metadata.mdx_n_fft_scale_set;
    this.hopLength = config.defaults.hopLength;
    this.dimF = config.metadata.mdx_dim_f_set;
@@ -1,4 +1,4 @@
-import type { LocalSeparatorModelConfig } from './types';
+import type { LocalSeparatorMdxModelConfig } from './types';
 import { LocalSeparatorCpuDsp } from './cpuDsp';
 import { reflectPad } from './shared';

@@ -140,7 +140,7 @@ export class LocalSeparatorGpuDsp {
  private paramBuffer: GPUBufferLike | null = null;
  private readonly pipeline: GPUComputePipelineLike;

-  public static async create(config: LocalSeparatorModelConfig): Promise<LocalSeparatorGpuDsp> {
+  public static async create(config: LocalSeparatorMdxModelConfig): Promise<LocalSeparatorGpuDsp> {
    if (!('gpu' in navigator)) {
      throw new Error('WebGPU is not available for GPU DSP.');
    }
@@ -164,7 +164,7 @@ export class LocalSeparatorGpuDsp {
    return new LocalSeparatorGpuDsp(config, device);
  }

-  private constructor(config: LocalSeparatorModelConfig, device: GPUDeviceLike) {
+  private constructor(config: LocalSeparatorMdxModelConfig, device: GPUDeviceLike) {
    this.device = device;
    this.cpuDsp = new LocalSeparatorCpuDsp(config);
    this.nFft = config.metadata.mdx_n_fft_scale_set;
@@ -1,7 +1,4 @@
-import {
-  LOCAL_SEPARATOR_MODEL_EXPECTED_SIZE_BYTES,
-  LOCAL_SEPARATOR_MODEL_FILENAME,
-} from './config';
+import type { LocalSeparatorModelConfig } from './types';
 import { OpfsModelCache, type ModelDownloadProgress } from '../opfsModelCache';

 const cache = new OpfsModelCache({ directoryName: 'models' });
@@ -9,34 +6,34 @@ const cache = new OpfsModelCache({ directoryName: 'models' });
 export { type ModelDownloadProgress };

 export class LocalSeparatorModelCache {
-  public static async exists(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<boolean> {
-    return cache.exists(filename, {
-      expectedSizeBytes: LOCAL_SEPARATOR_MODEL_EXPECTED_SIZE_BYTES,
+  public static async exists(modelConfig: LocalSeparatorModelConfig): Promise<boolean> {
+    return cache.exists(modelConfig.filename, {
+      expectedSizeBytes: modelConfig.download.expectedSizeBytes,
    });
  }

-  public static async getFile(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<File> {
-    return cache.getFile(filename);
+  public static async getFile(modelConfig: LocalSeparatorModelConfig): Promise<File> {
+    return cache.getFile(modelConfig.filename);
  }

-  public static async getArrayBuffer(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<ArrayBuffer> {
-    return cache.getArrayBuffer(filename);
+  public static async getArrayBuffer(modelConfig: LocalSeparatorModelConfig): Promise<ArrayBuffer> {
+    return cache.getArrayBuffer(modelConfig.filename);
  }

-  public static async delete(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<void> {
-    await cache.delete(filename);
+  public static async delete(modelConfig: LocalSeparatorModelConfig): Promise<void> {
+    await cache.delete(modelConfig.filename);
  }

  public static async download(
+    modelConfig: LocalSeparatorModelConfig,
    sourceUrl: string,
-    filename: string = LOCAL_SEPARATOR_MODEL_FILENAME,
    onProgress?: (progress: ModelDownloadProgress) => void,
  ): Promise<void> {
    await cache.download(
      sourceUrl,
-      filename,
+      modelConfig.filename,
      {
-        expectedSizeBytes: LOCAL_SEPARATOR_MODEL_EXPECTED_SIZE_BYTES,
+        expectedSizeBytes: modelConfig.download.expectedSizeBytes,
      },
      onProgress,
    );
@@ -1,10 +1,13 @@
 import * as ort from 'onnxruntime-web/webgpu';
+import { DemucsProcessor } from 'demucs-web';
 import { LocalSeparatorCpuDsp } from './cpuDsp';
 import { LocalSeparatorGpuDsp } from './gpuDsp';
 import { LocalSeparatorTimingCollector } from './timing';
 import type {
  LocalRuntimeProvider,
+  LocalSeparatorDemucsModelConfig,
  LocalSeparatorModelConfig,
+  LocalSeparatorMdxModelConfig,
  LocalSeparatorProgress,
  StereoChannels,
 } from './types';
@@ -76,8 +79,8 @@ function unpackBatchOutput(outputData: Float32Array, batchInfo: { dims: number[]
 class BrowserMdxSeparator {
  private readonly session: ort.InferenceSession;
  private readonly runtimeProvider: LocalRuntimeProvider;
-  private readonly defaults: LocalSeparatorModelConfig['defaults'];
-  private readonly metadata: LocalSeparatorModelConfig['metadata'];
+  private readonly defaults: LocalSeparatorMdxModelConfig['defaults'];
+  private readonly metadata: LocalSeparatorMdxModelConfig['metadata'];
  public onProgress: (progress: LocalSeparatorProgress) => void;
  private overlap: number;
  private runtimeBatchSize: number;
@@ -97,7 +100,7 @@ class BrowserMdxSeparator {
  public static async create(
    session: ort.InferenceSession,
    runtimeProvider: LocalRuntimeProvider,
-    config: LocalSeparatorModelConfig,
+    config: LocalSeparatorMdxModelConfig,
    options: BrowserMdxSeparatorOptions = {},
  ): Promise<BrowserMdxSeparator> {
    const timing = options.timing ?? new LocalSeparatorTimingCollector('mdx-separation');
@@ -139,7 +142,7 @@ class BrowserMdxSeparator {
  private constructor(
    session: ort.InferenceSession,
    runtimeProvider: LocalRuntimeProvider,
-    config: LocalSeparatorModelConfig,
+    config: LocalSeparatorMdxModelConfig,
    options: BrowserMdxSeparatorOptions & {
      dsp: LocalSeparatorCpuDsp | LocalSeparatorGpuDsp;
      dspMode: 'cpu' | 'gpu-hybrid';
@@ -470,10 +473,24 @@ function concatChannelPairs(chunks: StereoChannels[]): StereoChannels {
  return [concatFloat32(chunks.map(chunk => chunk[0])), concatFloat32(chunks.map(chunk => chunk[1]))];
 }

-export async function runLocalSeparator(options: {
+function getProviderLabel(runtimeProvider: LocalRuntimeProvider): string {
+  return runtimeProvider === 'webgpu' ? 'GPU/WebGPU' : 'CPU/wasm';
+}
+
+function getChunkDurationSeconds(
+  requestedChunkDurationSeconds: number | null,
+  modelConfig: LocalSeparatorModelConfig,
+): number | null {
+  if (requestedChunkDurationSeconds != null) {
+    return requestedChunkDurationSeconds;
+  }
+  return modelConfig.defaultChunkDurationSeconds;
+}
+
+async function runMdxLocalSeparator(options: {
  session: ort.InferenceSession;
  runtimeProvider: LocalRuntimeProvider;
-  modelConfig: LocalSeparatorModelConfig;
+  modelConfig: LocalSeparatorMdxModelConfig;
  audioBuffer: ArrayBuffer;
  chunkDurationSeconds: number | null;
  overlap: number;
@@ -486,7 +503,7 @@ export async function runLocalSeparator(options: {
 }> {
  const timing = new LocalSeparatorTimingCollector('local-separation');
  const decoded = await timing.measureAsync('decode', () => decodeAudioToStereo(options.audioBuffer));
-  log(`Running browser MDX separation on ${options.runtimeProvider === 'webgpu' ? 'GPU/WebGPU' : 'CPU/wasm'}...`);
+  log(`Running browser MDX separation on ${getProviderLabel(options.runtimeProvider)}...`);

  const separator = await BrowserMdxSeparator.create(
    options.session,
@@ -502,11 +519,11 @@ export async function runLocalSeparator(options: {
  );

  try {
-    const outputs = await separateWithOptionalChunking(
+    const outputs = await separateMdxWithOptionalChunking(
      separator,
      decoded,
      timing,
-      options.chunkDurationSeconds,
+      getChunkDurationSeconds(options.chunkDurationSeconds, options.modelConfig),
      options.modelConfig,
      options.onProgress,
    );
@@ -518,7 +535,7 @@ export async function runLocalSeparator(options: {
        { name: outputs.primaryStem, blob: primaryBlob },
        { name: outputs.secondaryStem, blob: secondaryBlob },
      ],
-      providerLabel: options.runtimeProvider === 'webgpu' ? 'GPU/WebGPU' : 'CPU/wasm',
+      providerLabel: getProviderLabel(options.runtimeProvider),
      debugSummary: separator.getDebugSummary({ model: options.modelConfig.filename }),
    };
  } finally {
@@ -527,12 +544,12 @@ export async function runLocalSeparator(options: {
  }
 }

-async function separateWithOptionalChunking(
+async function separateMdxWithOptionalChunking(
  separator: BrowserMdxSeparator,
  decoded: StereoChannels,
  timing: LocalSeparatorTimingCollector,
  chunkDurationSeconds: number | null,
-  modelConfig: LocalSeparatorModelConfig,
+  modelConfig: LocalSeparatorMdxModelConfig,
  onProgress: (progress: LocalSeparatorProgress) => void,
 ): Promise<{
  stems: Record<string, StereoChannels>;
@@ -594,3 +611,171 @@ async function separateWithOptionalChunking(
    secondaryStem,
  };
 }
+
+type DemucsStemName = 'Vocals' | 'Drums' | 'Bass' | 'Others';
+
+function toDemucsStemMap(result: Awaited<ReturnType<DemucsProcessor['separate']>>): Record<DemucsStemName, StereoChannels> {
+  return {
+    Vocals: [result.vocals.left, result.vocals.right],
+    Drums: [result.drums.left, result.drums.right],
+    Bass: [result.bass.left, result.bass.right],
+    Others: [result.other.left, result.other.right],
+  };
+}
+
+async function runDemucsLocalSeparator(options: {
+  session: ort.InferenceSession;
+  runtimeProvider: LocalRuntimeProvider;
+  modelConfig: LocalSeparatorDemucsModelConfig;
+  audioBuffer: ArrayBuffer;
+  chunkDurationSeconds: number | null;
+  onProgress: (progress: LocalSeparatorProgress) => void;
+}): Promise<{
+  stems: Array<{ name: string; blob: Blob }>;
+  providerLabel: string;
+  debugSummary: Record<string, unknown>;
+}> {
+  const timing = new LocalSeparatorTimingCollector('local-demucs-separation');
+  const decoded = await timing.measureAsync('decode', () => decodeAudioToStereo(options.audioBuffer));
+  log(`Running browser Demucs separation on ${getProviderLabel(options.runtimeProvider)}...`);
+
+  const processor = new DemucsProcessor({
+    ort,
+    onProgress: ({ progress, currentSegment, totalSegments }) => {
+      options.onProgress({
+        stage: 'main',
+        passLabel: 'Demucs pass',
+        percent: progress * 100,
+        processedChunks: currentSegment,
+        totalChunks: totalSegments,
+      });
+    },
+    onLog: (phase, message) => log(`demucs:${phase}`, message),
+  });
+  processor.session = options.session;
+
+  const outputMap = await separateDemucsWithOptionalChunking(
+    processor,
+    decoded,
+    timing,
+    getChunkDurationSeconds(options.chunkDurationSeconds, options.modelConfig),
+    options.modelConfig,
+    options.onProgress,
+  );
+
+  const stems = options.modelConfig.outputStemNames.map(name => ({
+    name,
+    blob: timing.measureSync(`wavEncode:${name}`, () => channelsToWavBlob(outputMap[name as DemucsStemName])),
+  }));
+
+  return {
+    stems,
+    providerLabel: getProviderLabel(options.runtimeProvider),
+    debugSummary: timing.getSummary({
+      model: options.modelConfig.filename,
+      runtimeProvider: options.runtimeProvider,
+      stemCount: stems.length,
+    }),
+  };
+}
+
+async function separateDemucsWithOptionalChunking(
+  processor: DemucsProcessor,
+  decoded: StereoChannels,
+  timing: LocalSeparatorTimingCollector,
+  chunkDurationSeconds: number | null,
+  modelConfig: LocalSeparatorDemucsModelConfig,
+  onProgress: (progress: LocalSeparatorProgress) => void,
+): Promise<Record<DemucsStemName, StereoChannels>> {
+  if (!chunkDurationSeconds) {
+    return timing.measureAsync('demucsSeparate', async () => toDemucsStemMap(await processor.separate(decoded[0], decoded[1])));
+  }
+
+  const chunkSamples = Math.max(1, Math.floor(chunkDurationSeconds * SAMPLE_RATE));
+  if (decoded[0].length <= chunkSamples) {
+    return timing.measureAsync('demucsSeparate', async () => toDemucsStemMap(await processor.separate(decoded[0], decoded[1])));
+  }
+
+  const totalChunks = Math.ceil(decoded[0].length / chunkSamples);
+  const chunkedOutputs: Record<DemucsStemName, StereoChannels[]> = {
+    Vocals: [],
+    Drums: [],
+    Bass: [],
+    Others: [],
+  };
+
+  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);
+
+    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 chunkResult = await timing.measureAsync(
+      'demucsChunkedSeparate',
+      async () => toDemucsStemMap(await processor.separate(chunk[0], chunk[1])),
+    );
+
+    (Object.keys(chunkedOutputs) as DemucsStemName[]).forEach(stemName => {
+      chunkedOutputs[stemName].push(chunkResult[stemName]);
+    });
+
+    onProgress({
+      stage: 'chunk-complete',
+      passLabel: `Audio chunk ${index + 1}/${totalChunks}: Demucs pass`,
+      percent: ((index + 1) / totalChunks) * 100,
+      processedChunks: index + 1,
+      totalChunks,
+    });
+  }
+
+  return {
+    Vocals: concatChannelPairs(chunkedOutputs.Vocals),
+    Drums: concatChannelPairs(chunkedOutputs.Drums),
+    Bass: concatChannelPairs(chunkedOutputs.Bass),
+    Others: concatChannelPairs(chunkedOutputs.Others),
+  };
+}
+
+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<string, unknown>;
+}> {
+  if (options.modelConfig.kind === 'demucs') {
+    return runDemucsLocalSeparator({
+      session: options.session,
+      runtimeProvider: options.runtimeProvider,
+      modelConfig: options.modelConfig,
+      audioBuffer: options.audioBuffer,
+      chunkDurationSeconds: options.chunkDurationSeconds,
+      onProgress: options.onProgress,
+    });
+  }
+
+  return runMdxLocalSeparator({
+    session: options.session,
+    runtimeProvider: options.runtimeProvider,
+    modelConfig: options.modelConfig,
+    audioBuffer: options.audioBuffer,
+    chunkDurationSeconds: options.chunkDurationSeconds,
+    overlap: options.overlap,
+    onProgress: options.onProgress,
+    onProviderChange: options.onProviderChange,
+  });
+}
@@ -1,5 +1,8 @@
 export type StereoChannels = [Float32Array, Float32Array];

+export type LocalSeparatorModelId = 'UVR-MDX-NET-Inst_HQ_3.onnx' | 'htdemucs_4s.onnx';
+export type LocalSeparatorModelKind = 'mdx' | 'demucs';
+
 export interface LocalSeparatorModelDefaults {
  sampleRate: number;
  hopLength: number;
@@ -13,7 +16,7 @@ export interface LocalSeparatorModelDefaults {
  matchMixOverlap: number;
 }

-export interface LocalSeparatorModelMetadata {
+export interface LocalSeparatorMdxMetadata {
  compensate: number;
  mdx_dim_f_set: number;
  mdx_dim_t_set: number;
@@ -21,14 +24,36 @@ export interface LocalSeparatorModelMetadata {
  primary_stem: string;
 }

-export interface LocalSeparatorModelConfig {
+export interface LocalSeparatorModelDownloadConfig {
+  configKey: 'general.uvr5_web_runtime.mdx_net_model_url' | 'general.uvr5_web_runtime.htdemucs_4s_model_url';
+  defaultUrl: string;
+  expectedSizeBytes: number;
+}
+
+interface LocalSeparatorModelConfigBase {
+  id: LocalSeparatorModelId;
  filename: string;
+  kind: LocalSeparatorModelKind;
  displayName: string;
  status: 'ready';
+  outputStemNames: string[];
+  defaultChunkDurationSeconds: number | null;
  defaults: LocalSeparatorModelDefaults;
-  metadata: LocalSeparatorModelMetadata;
+  download: LocalSeparatorModelDownloadConfig;
 }

+export interface LocalSeparatorMdxModelConfig extends LocalSeparatorModelConfigBase {
+  kind: 'mdx';
+  metadata: LocalSeparatorMdxMetadata;
+}
+
+export interface LocalSeparatorDemucsModelConfig extends LocalSeparatorModelConfigBase {
+  kind: 'demucs';
+  metadata: null;
+}
+
+export type LocalSeparatorModelConfig = LocalSeparatorMdxModelConfig | LocalSeparatorDemucsModelConfig;
+
 export interface LocalSeparatorProgress {
  stage: string;
  passLabel: string;