459 lines
17 KiB
JavaScript
459 lines
17 KiB
JavaScript
/**
|
|
* Demucs Web - Stem Extraction for SunoAce
|
|
*/
|
|
import * as ort from 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.21.0/dist/ort.all.mjs';
|
|
import { DemucsProcessor, CONSTANTS } from './src/index.js';
|
|
|
|
const { SAMPLE_RATE, TRAINING_SAMPLES, TRACKS, DEFAULT_MODEL_URL } = CONSTANTS;
|
|
|
|
const LOCAL_MODEL_URL = '../models/htdemucs_embedded.onnx';
|
|
|
|
let processor = null;
|
|
let audioContext = null;
|
|
let audioBuffer = null;
|
|
let isProcessing = false;
|
|
|
|
// DOM elements
|
|
const dropZone = document.getElementById('dropZone');
|
|
const fileInput = document.getElementById('fileInput');
|
|
const processBtn = document.getElementById('processBtn');
|
|
const progressFill = document.getElementById('progressFill');
|
|
const status = document.getElementById('status');
|
|
const results = document.getElementById('results');
|
|
const trackList = document.getElementById('trackList');
|
|
const backendBadge = document.getElementById('backendBadge');
|
|
const audioFileName = document.getElementById('audioFileName');
|
|
const statusDetail = document.getElementById('statusDetail');
|
|
const statsRow = document.getElementById('statsRow');
|
|
const statElapsed = document.getElementById('statElapsed');
|
|
const statSegment = document.getElementById('statSegment');
|
|
const statSpeed = document.getElementById('statSpeed');
|
|
const statETA = document.getElementById('statETA');
|
|
|
|
let processStartTime = null;
|
|
|
|
function log(phase, message) {
|
|
const now = new Date();
|
|
const timeStr = now.toLocaleTimeString('en-US', { hour12: false });
|
|
const logLine = document.createElement('div');
|
|
logLine.className = 'text-zinc-400 py-1 border-b border-zinc-800/50 last:border-0';
|
|
logLine.innerHTML = `<span class="text-emerald-400">[${timeStr}]</span> <span class="text-teal-400">[${phase}]</span> ${message}`;
|
|
statusDetail.appendChild(logLine);
|
|
statusDetail.scrollTop = statusDetail.scrollHeight;
|
|
console.log(`[${phase}] ${message}`);
|
|
}
|
|
|
|
function formatTime(seconds) {
|
|
if (!isFinite(seconds) || seconds < 0) return '--:--';
|
|
const mins = Math.floor(seconds / 60);
|
|
const secs = Math.floor(seconds % 60);
|
|
return `${mins}:${secs.toString().padStart(2, '0')}`;
|
|
}
|
|
|
|
async function init() {
|
|
let backend = 'wasm';
|
|
|
|
if ('gpu' in navigator) {
|
|
try {
|
|
const gpuAdapter = await navigator.gpu.requestAdapter();
|
|
if (gpuAdapter) {
|
|
backend = 'webgpu';
|
|
}
|
|
} catch (e) {
|
|
console.log('WebGPU not available:', e);
|
|
}
|
|
}
|
|
|
|
ort.env.wasm.numThreads = navigator.hardwareConcurrency || 4;
|
|
|
|
if (backend === 'webgpu') {
|
|
ort.env.webgpu = ort.env.webgpu || {};
|
|
ort.env.webgpu.powerPreference = 'high-performance';
|
|
backendBadge.textContent = 'WebGPU (GPU)';
|
|
backendBadge.className = 'badge badge-gpu';
|
|
} else {
|
|
const threads = navigator.hardwareConcurrency || 4;
|
|
backendBadge.textContent = `WASM (${threads} threads)`;
|
|
backendBadge.className = 'badge badge-cpu';
|
|
}
|
|
|
|
processor = new DemucsProcessor({
|
|
ort,
|
|
onProgress: ({ progress, currentSegment, totalSegments }) => {
|
|
progressFill.style.width = (5 + progress * 90) + '%';
|
|
|
|
const elapsed = (Date.now() - processStartTime) / 1000;
|
|
statElapsed.textContent = formatTime(elapsed);
|
|
statSegment.textContent = `${currentSegment}/${totalSegments}`;
|
|
|
|
if (currentSegment > 0 && audioBuffer) {
|
|
const processedDuration = (currentSegment / totalSegments) * audioBuffer.duration;
|
|
const speed = processedDuration / elapsed;
|
|
statSpeed.textContent = speed.toFixed(2) + 'x';
|
|
|
|
const remainingSegments = totalSegments - currentSegment;
|
|
const avgTimePerSegment = elapsed / currentSegment;
|
|
const eta = remainingSegments * avgTimePerSegment;
|
|
statETA.textContent = formatTime(eta);
|
|
}
|
|
},
|
|
onLog: log,
|
|
onDownloadProgress: (loaded, total) => {
|
|
const percent = ((loaded / total) * 100).toFixed(1);
|
|
const loadedMB = (loaded / 1024 / 1024).toFixed(1);
|
|
const totalMB = (total / 1024 / 1024).toFixed(1);
|
|
status.textContent = `Downloading model... ${loadedMB}MB / ${totalMB}MB (${percent}%)`;
|
|
progressFill.style.width = (loaded / total * 100) + '%';
|
|
}
|
|
});
|
|
|
|
status.textContent = 'Loading AI model...';
|
|
|
|
try {
|
|
try {
|
|
status.textContent = 'Downloading model (~172MB)...';
|
|
await processor.loadModel(DEFAULT_MODEL_URL);
|
|
} catch {
|
|
status.textContent = 'Loading local model...';
|
|
await processor.loadModel(LOCAL_MODEL_URL);
|
|
}
|
|
status.textContent = 'Ready - Select an audio file';
|
|
progressFill.style.width = '0%';
|
|
} catch (e) {
|
|
status.textContent = 'Failed to load model: ' + e.message;
|
|
console.error('Failed to load model:', e);
|
|
}
|
|
|
|
audioContext = new (window.AudioContext || window.webkitAudioContext)({
|
|
sampleRate: SAMPLE_RATE
|
|
});
|
|
|
|
// Check for audio URL parameter and auto-start
|
|
const urlParams = new URLSearchParams(window.location.search);
|
|
const audioUrl = urlParams.get('audioUrl');
|
|
if (audioUrl) {
|
|
await loadAudioFromUrl(audioUrl);
|
|
}
|
|
}
|
|
|
|
async function loadAudioFromUrl(url) {
|
|
try {
|
|
status.textContent = 'Loading audio...';
|
|
const fileName = decodeURIComponent(url.split('/').pop() || 'audio.mp3');
|
|
audioFileName.textContent = fileName;
|
|
|
|
// Force fresh fetch to avoid 304 Not Modified with empty body
|
|
const response = await fetch(url, { cache: 'no-store' });
|
|
if (!response.ok) throw new Error(`HTTP ${response.status}`);
|
|
|
|
const arrayBuffer = await response.arrayBuffer();
|
|
audioBuffer = await audioContext.decodeAudioData(arrayBuffer);
|
|
|
|
const duration = audioBuffer.duration.toFixed(1);
|
|
status.textContent = `Loaded: ${duration}s - Starting extraction...`;
|
|
processBtn.disabled = false;
|
|
|
|
// Auto-start extraction
|
|
setTimeout(() => startProcessing(), 500);
|
|
} catch (e) {
|
|
status.textContent = 'Failed to load audio: ' + e.message;
|
|
console.error('Failed to load audio from URL:', e);
|
|
}
|
|
}
|
|
|
|
// Drag and drop handlers
|
|
dropZone.addEventListener('click', () => fileInput.click());
|
|
dropZone.addEventListener('dragover', (e) => {
|
|
e.preventDefault();
|
|
dropZone.classList.add('border-emerald-500', 'bg-emerald-500/5');
|
|
});
|
|
dropZone.addEventListener('dragleave', () => {
|
|
dropZone.classList.remove('border-emerald-500', 'bg-emerald-500/5');
|
|
});
|
|
dropZone.addEventListener('drop', (e) => {
|
|
e.preventDefault();
|
|
dropZone.classList.remove('border-emerald-500', 'bg-emerald-500/5');
|
|
const file = e.dataTransfer.files[0];
|
|
if (file && file.type.startsWith('audio/')) {
|
|
handleFile(file);
|
|
}
|
|
});
|
|
fileInput.addEventListener('change', (e) => {
|
|
const file = e.target.files[0];
|
|
if (file) handleFile(file);
|
|
});
|
|
|
|
async function handleFile(file) {
|
|
audioFileName.textContent = file.name;
|
|
status.textContent = 'Reading audio...';
|
|
|
|
try {
|
|
const arrayBuffer = await file.arrayBuffer();
|
|
audioBuffer = await audioContext.decodeAudioData(arrayBuffer);
|
|
const duration = audioBuffer.duration.toFixed(1);
|
|
status.textContent = `Loaded: ${duration}s - Ready to extract`;
|
|
processBtn.disabled = false;
|
|
} catch (e) {
|
|
status.textContent = 'Failed to read audio: ' + e.message;
|
|
console.error('Failed to decode audio:', e);
|
|
}
|
|
}
|
|
|
|
processBtn.addEventListener('click', startProcessing);
|
|
|
|
async function startProcessing() {
|
|
if (!audioBuffer || !processor || isProcessing) return;
|
|
|
|
isProcessing = true;
|
|
processBtn.disabled = true;
|
|
processBtn.textContent = 'Processing...';
|
|
results.classList.remove('visible');
|
|
processStartTime = Date.now();
|
|
statusDetail.innerHTML = '';
|
|
statusDetail.classList.add('visible');
|
|
statsRow.classList.add('visible');
|
|
|
|
try {
|
|
log('Init', 'Starting stem extraction...');
|
|
status.textContent = 'Preparing audio...';
|
|
progressFill.style.width = '2%';
|
|
|
|
let leftChannel = audioBuffer.getChannelData(0);
|
|
let rightChannel = audioBuffer.numberOfChannels > 1
|
|
? audioBuffer.getChannelData(1)
|
|
: leftChannel;
|
|
|
|
if (audioBuffer.sampleRate !== SAMPLE_RATE) {
|
|
log('Resample', `${audioBuffer.sampleRate}Hz → ${SAMPLE_RATE}Hz`);
|
|
const ratio = SAMPLE_RATE / audioBuffer.sampleRate;
|
|
const newLength = Math.floor(leftChannel.length * ratio);
|
|
const newLeft = new Float32Array(newLength);
|
|
const newRight = new Float32Array(newLength);
|
|
|
|
for (let i = 0; i < newLength; i++) {
|
|
const srcIdx = i / ratio;
|
|
const idx0 = Math.floor(srcIdx);
|
|
const idx1 = Math.min(idx0 + 1, leftChannel.length - 1);
|
|
const frac = srcIdx - idx0;
|
|
newLeft[i] = leftChannel[idx0] * (1 - frac) + leftChannel[idx1] * frac;
|
|
newRight[i] = rightChannel[idx0] * (1 - frac) + rightChannel[idx1] * frac;
|
|
}
|
|
|
|
leftChannel = newLeft;
|
|
rightChannel = newRight;
|
|
}
|
|
|
|
status.textContent = 'Extracting stems...';
|
|
const separatedTracks = await processor.separate(leftChannel, rightChannel);
|
|
displayResults(separatedTracks);
|
|
|
|
const totalTime = ((Date.now() - processStartTime) / 1000).toFixed(1);
|
|
const speedRatio = (audioBuffer.duration / parseFloat(totalTime)).toFixed(2);
|
|
|
|
log('Done', `Completed in ${totalTime}s (${speedRatio}x realtime)`);
|
|
status.textContent = `Complete! Extracted 4 stems in ${totalTime}s`;
|
|
progressFill.style.width = '100%';
|
|
|
|
} catch (e) {
|
|
status.textContent = 'Processing failed: ' + e.message;
|
|
console.error('Processing failed:', e);
|
|
}
|
|
|
|
isProcessing = false;
|
|
processBtn.disabled = false;
|
|
processBtn.textContent = 'Extract Stems';
|
|
}
|
|
|
|
// Store track URLs for download all feature
|
|
let trackUrls = {};
|
|
|
|
function displayResults(tracks) {
|
|
trackList.innerHTML = '';
|
|
trackUrls = {};
|
|
|
|
const TRACK_CONFIG = {
|
|
drums: { icon: '🥁', label: 'Drums' },
|
|
bass: { icon: '🎸', label: 'Bass' },
|
|
other: { icon: '🎹', label: 'Instrumental' },
|
|
vocals: { icon: '🎤', label: 'Vocals' }
|
|
};
|
|
|
|
for (const [name, track] of Object.entries(tracks)) {
|
|
const config = TRACK_CONFIG[name] || { icon: '🎵', label: name };
|
|
const trackBuffer = audioContext.createBuffer(2, track.left.length, SAMPLE_RATE);
|
|
trackBuffer.getChannelData(0).set(track.left);
|
|
trackBuffer.getChannelData(1).set(track.right);
|
|
|
|
const audioBlob = audioBufferToWav(trackBuffer);
|
|
const audioUrl = URL.createObjectURL(audioBlob);
|
|
const trackId = `track-${name}`;
|
|
const fileName = config.label.toLowerCase();
|
|
|
|
// Store for download all
|
|
trackUrls[fileName] = audioUrl;
|
|
|
|
const trackDiv = document.createElement('div');
|
|
trackDiv.className = 'track';
|
|
trackDiv.innerHTML = `
|
|
<div class="track-row">
|
|
<div class="track-info">
|
|
<div class="track-icon ${name}">${config.icon}</div>
|
|
<div>
|
|
<div class="track-name">${config.label}</div>
|
|
<div class="track-duration">${formatTime(trackBuffer.duration)}</div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="track-player">
|
|
<button id="play-${trackId}" class="play-btn" onclick="togglePlay('${trackId}')">
|
|
<svg fill="currentColor" viewBox="0 0 24 24"><path d="M8 5v14l11-7z"/></svg>
|
|
</button>
|
|
|
|
<div id="progress-bg-${trackId}" class="track-progress" onclick="seekTrack(event, '${trackId}')">
|
|
<div id="progress-${trackId}" class="track-progress-fill ${name}"></div>
|
|
</div>
|
|
|
|
<span id="time-${trackId}" class="track-time">0:00 / ${formatTime(trackBuffer.duration)}</span>
|
|
</div>
|
|
|
|
<a href="${audioUrl}" download="${fileName}.wav" class="download-btn">
|
|
<svg fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
|
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"/>
|
|
</svg>
|
|
WAV
|
|
</a>
|
|
</div>
|
|
|
|
<audio id="audio-${trackId}" src="${audioUrl}" preload="metadata"></audio>
|
|
`;
|
|
|
|
trackList.appendChild(trackDiv);
|
|
|
|
const audio = document.getElementById(`audio-${trackId}`);
|
|
audio.addEventListener('timeupdate', () => updateProgress(trackId, audio));
|
|
audio.addEventListener('ended', () => resetPlayer(trackId));
|
|
}
|
|
|
|
results.classList.add('visible');
|
|
}
|
|
|
|
// Download all stems
|
|
window.downloadAllStems = function() {
|
|
const entries = Object.entries(trackUrls);
|
|
let index = 0;
|
|
|
|
function downloadNext() {
|
|
if (index >= entries.length) return;
|
|
const [name, url] = entries[index];
|
|
const a = document.createElement('a');
|
|
a.href = url;
|
|
a.download = `${name}.wav`;
|
|
document.body.appendChild(a);
|
|
a.click();
|
|
document.body.removeChild(a);
|
|
index++;
|
|
setTimeout(downloadNext, 500);
|
|
}
|
|
|
|
downloadNext();
|
|
};
|
|
|
|
// Player functions (global scope for onclick handlers)
|
|
window.togglePlay = function(trackId) {
|
|
const audio = document.getElementById(`audio-${trackId}`);
|
|
const playBtn = document.getElementById(`play-${trackId}`);
|
|
|
|
// Pause all other tracks
|
|
document.querySelectorAll('audio').forEach(a => {
|
|
if (a.id !== `audio-${trackId}` && !a.paused) {
|
|
a.pause();
|
|
const otherId = a.id.replace('audio-', '');
|
|
resetPlayer(otherId);
|
|
}
|
|
});
|
|
|
|
if (audio.paused) {
|
|
audio.play();
|
|
playBtn.innerHTML = `<svg class="w-4 h-4" fill="currentColor" viewBox="0 0 24 24"><path d="M6 4h4v16H6zm8 0h4v16h-4z"/></svg>`;
|
|
} else {
|
|
audio.pause();
|
|
playBtn.innerHTML = `<svg class="w-4 h-4 ml-0.5" fill="currentColor" viewBox="0 0 24 24"><path d="M8 5v14l11-7z"/></svg>`;
|
|
}
|
|
};
|
|
|
|
window.seekTrack = function(event, trackId) {
|
|
const audio = document.getElementById(`audio-${trackId}`);
|
|
const progressBg = document.getElementById(`progress-bg-${trackId}`);
|
|
const rect = progressBg.getBoundingClientRect();
|
|
const percent = (event.clientX - rect.left) / rect.width;
|
|
audio.currentTime = percent * audio.duration;
|
|
};
|
|
|
|
function updateProgress(trackId, audio) {
|
|
const progress = document.getElementById(`progress-${trackId}`);
|
|
const timeDisplay = document.getElementById(`time-${trackId}`);
|
|
const percent = (audio.currentTime / audio.duration) * 100;
|
|
progress.style.width = `${percent}%`;
|
|
timeDisplay.textContent = `${formatTime(audio.currentTime)} / ${formatTime(audio.duration)}`;
|
|
}
|
|
|
|
function resetPlayer(trackId) {
|
|
const playBtn = document.getElementById(`play-${trackId}`);
|
|
const progress = document.getElementById(`progress-${trackId}`);
|
|
playBtn.innerHTML = `<svg class="w-4 h-4 ml-0.5" fill="currentColor" viewBox="0 0 24 24"><path d="M8 5v14l11-7z"/></svg>`;
|
|
progress.style.width = '0%';
|
|
}
|
|
|
|
function audioBufferToWav(buffer) {
|
|
const numChannels = buffer.numberOfChannels;
|
|
const sampleRate = buffer.sampleRate;
|
|
const bitDepth = 16;
|
|
const bytesPerSample = bitDepth / 8;
|
|
const blockAlign = numChannels * bytesPerSample;
|
|
const samples = buffer.length;
|
|
const dataSize = samples * blockAlign;
|
|
const bufferSize = 44 + dataSize;
|
|
|
|
const arrayBuffer = new ArrayBuffer(bufferSize);
|
|
const view = new DataView(arrayBuffer);
|
|
|
|
const writeString = (offset, string) => {
|
|
for (let i = 0; i < string.length; i++) {
|
|
view.setUint8(offset + i, string.charCodeAt(i));
|
|
}
|
|
};
|
|
|
|
writeString(0, 'RIFF');
|
|
view.setUint32(4, bufferSize - 8, true);
|
|
writeString(8, 'WAVE');
|
|
writeString(12, 'fmt ');
|
|
view.setUint32(16, 16, true);
|
|
view.setUint16(20, 1, true);
|
|
view.setUint16(22, numChannels, true);
|
|
view.setUint32(24, sampleRate, true);
|
|
view.setUint32(28, sampleRate * blockAlign, true);
|
|
view.setUint16(32, blockAlign, true);
|
|
view.setUint16(34, bitDepth, true);
|
|
writeString(36, 'data');
|
|
view.setUint32(40, dataSize, true);
|
|
|
|
const channels = [];
|
|
for (let c = 0; c < numChannels; c++) {
|
|
channels.push(buffer.getChannelData(c));
|
|
}
|
|
|
|
let offset = 44;
|
|
for (let i = 0; i < samples; i++) {
|
|
for (let c = 0; c < numChannels; c++) {
|
|
const sample = Math.max(-1, Math.min(1, channels[c][i]));
|
|
const intSample = sample < 0 ? sample * 0x8000 : sample * 0x7FFF;
|
|
view.setInt16(offset, intSample, true);
|
|
offset += 2;
|
|
}
|
|
}
|
|
|
|
return new Blob([arrayBuffer], { type: 'audio/wav' });
|
|
}
|
|
|
|
init();
|