Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import spaces | |
| PAGE = r""" | |
| <div id="sam-browser-app" class="sam-app"> | |
| <style> | |
| #sam-browser-app{--sam-accent:#7c6cf2;--sam-accent2:#a89dff;max-width:1120px;margin:0 auto;font-family:Inter,ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;color:var(--body-text-color)} | |
| #sam-browser-app *{box-sizing:border-box} | |
| #sam-browser-app .hero{padding:26px 0 20px} | |
| #sam-browser-app .eyebrow{font-size:12px;font-weight:800;letter-spacing:.13em;text-transform:uppercase;color:var(--sam-accent2);margin-bottom:9px} | |
| #sam-browser-app h1{font-size:clamp(28px,4vw,40px);letter-spacing:-.04em;line-height:1.12;margin:0 0 12px;font-weight:780} | |
| #sam-browser-app h2{font-size:18px;letter-spacing:-.02em;margin:0;font-weight:750} | |
| #sam-browser-app p{line-height:1.55} | |
| #sam-browser-app .subhead{font-size:15px;max-width:760px;color:var(--body-text-color);opacity:.76;margin:0} | |
| #sam-browser-app .pill{display:inline-flex;align-items:center;gap:7px;border:1px solid var(--border-color-primary);border-radius:999px;padding:7px 11px;font-size:12px;font-weight:650;margin-top:16px} | |
| #sam-browser-app .dot{width:7px;height:7px;border-radius:50%;background:#3dbb86;display:inline-block} | |
| #sam-browser-app .panel{border:1px solid var(--border-color-primary);border-radius:18px;padding:clamp(16px,3vw,24px);margin:15px 0;background:var(--background-fill-primary);box-shadow:0 5px 22px rgba(0,0,0,.035)} | |
| #sam-browser-app .section-head{display:flex;align-items:center;gap:12px;margin-bottom:19px} | |
| #sam-browser-app .step{display:grid;place-items:center;flex:0 0 34px;height:34px;border-radius:11px;background:color-mix(in srgb,var(--sam-accent) 14%,transparent);color:var(--sam-accent2);font-weight:800} | |
| #sam-browser-app .muted{opacity:.68;font-size:13px;line-height:1.5} | |
| #sam-browser-app label{display:block;margin:0 0 7px;font-size:13px;font-weight:700} | |
| #sam-browser-app input[type=file],#sam-browser-app input[type=text],#sam-browser-app input[type=number],#sam-browser-app select{width:100%;min-width:0;padding:12px 13px;border-radius:10px;border:1px solid var(--border-color-primary);background:var(--input-background-fill);color:var(--body-text-color);font:inherit;font-size:14px} | |
| #sam-browser-app input[type=file]{padding:10px} | |
| #sam-browser-app .file-wrap{padding:16px;border:1px dashed var(--border-color-primary);border-radius:14px;background:var(--background-fill-secondary)} | |
| #sam-browser-app .controls{display:grid;grid-template-columns:minmax(0,2fr) minmax(150px,1fr) minmax(110px,.7fr);gap:14px;margin-top:18px} | |
| #sam-browser-app .field{min-width:0} | |
| #sam-browser-app .actions{display:flex;align-items:center;gap:12px;flex-wrap:wrap;margin-top:20px} | |
| #sam-browser-app button{padding:12px 17px;border:0;border-radius:11px;background:var(--sam-accent);color:white;cursor:pointer;font-weight:750;font-size:14px;transition:filter .15s,transform .15s} | |
| #sam-browser-app button:hover:not(:disabled){filter:brightness(1.08);transform:translateY(-1px)} | |
| #sam-browser-app button:disabled{opacity:.55;cursor:wait;transform:none} | |
| #sam-browser-app .secondary{background:var(--button-secondary-background-fill);color:var(--button-secondary-text-color);border:1px solid var(--border-color-primary)} | |
| #sam-browser-app .status{display:flex;align-items:center;gap:9px;min-height:42px;padding:10px 13px;border-radius:11px;background:var(--background-fill-secondary);font-size:13px;flex:1;min-width:180px} | |
| #sam-browser-app .status-dot{width:8px;height:8px;flex:0 0 8px;border-radius:50%;background:#a1a1aa} | |
| #sam-browser-app .status[data-state=busy] .status-dot{background:#e9a23b;box-shadow:0 0 0 4px color-mix(in srgb,#e9a23b 17%,transparent)} | |
| #sam-browser-app .status[data-state=success] .status-dot{background:#35b67a} | |
| #sam-browser-app .status[data-state=error] .status-dot{background:#e05a68} | |
| #sam-browser-app .audio-grid{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:14px} | |
| #sam-browser-app .audio-card{min-width:0;border:1px solid var(--border-color-primary);border-radius:15px;padding:15px;background:var(--background-fill-secondary)} | |
| #sam-browser-app .audio-card-head{display:flex;align-items:center;justify-content:space-between;gap:8px;margin-bottom:13px} | |
| #sam-browser-app .audio-title{font-size:14px;font-weight:800;letter-spacing:-.01em} | |
| #sam-browser-app .tag{font-size:10px;letter-spacing:.08em;text-transform:uppercase;font-weight:800;border-radius:6px;padding:5px 7px;background:color-mix(in srgb,var(--sam-accent) 15%,transparent);color:var(--sam-accent2)} | |
| #sam-browser-app .tag.neutral{background:var(--border-color-primary);color:var(--body-text-color)} | |
| #sam-browser-app .wave-wrap{position:relative;width:100%;height:100px;border-radius:9px;overflow:hidden;background:color-mix(in srgb,var(--background-fill-primary) 78%,transparent)} | |
| #sam-browser-app canvas{display:block;width:100%;height:100px;cursor:pointer;touch-action:pan-y} | |
| #sam-browser-app .wave-empty[hidden]{display:none} | |
| #sam-browser-app .wave-empty{position:absolute;inset:0;display:grid;place-items:center;text-align:center;padding:10px;color:var(--body-text-color);opacity:.45;font-size:12px;pointer-events:none} | |
| #sam-browser-app .audio-meta{display:flex;justify-content:space-between;gap:8px;font-size:11px;opacity:.62;margin:9px 0 7px;font-variant-numeric:tabular-nums} | |
| #sam-browser-app audio{display:block;width:100%;height:40px;margin:7px 0 11px} | |
| #sam-browser-app .download{display:inline-flex;align-items:center;gap:7px;text-decoration:none;font-size:12px;font-weight:750;color:var(--sam-accent2);padding:7px 0} | |
| #sam-browser-app .download[hidden]{display:none} | |
| #sam-browser-app .empty-results{padding:24px 14px;border:1px dashed var(--border-color-primary);border-radius:14px;text-align:center;color:var(--body-text-color);opacity:.65;font-size:13px} | |
| #sam-browser-app .privacy{display:flex;gap:10px;align-items:flex-start;padding:12px 14px;border-radius:11px;background:color-mix(in srgb,#35b67a 8%,transparent);font-size:12px;line-height:1.5;margin-top:16px} | |
| #sam-browser-app .privacy svg{flex:0 0 17px;margin-top:1px} | |
| #sam-browser-app .footer-note{font-size:12px;opacity:.58;text-align:center;margin:18px 0 8px;line-height:1.5} | |
| #sam-browser-app .credits{font-size:12px;opacity:.7;text-align:center;margin:0 0 28px;line-height:1.8} | |
| #sam-browser-app .credits a{color:var(--sam-accent2);text-decoration:none;font-weight:650} | |
| #sam-browser-app .credits a:hover{text-decoration:underline} | |
| @media(max-width:820px){#sam-browser-app .audio-grid{grid-template-columns:1fr}#sam-browser-app .controls{grid-template-columns:1fr 1fr}#sam-browser-app .controls .prompt-field{grid-column:1/-1}} | |
| @media(max-width:480px){#sam-browser-app .controls{grid-template-columns:1fr}#sam-browser-app .controls .prompt-field{grid-column:auto}#sam-browser-app .actions{align-items:stretch}#sam-browser-app .actions button{width:100%}#sam-browser-app .status{flex-basis:100%}} | |
| </style> | |
| <header class="hero"> | |
| <div class="eyebrow">Audio separation studio</div> | |
| <h1>SAM Audio <span style="color:var(--sam-accent2)">Studio</span></h1> | |
| <p class="subhead">Separate speech, music, or other sounds from a recording using Meta's SAM Audio, running locally right in your browser.<br><br>Once downloaded, the model takes up about 2.7 GB in the browser cache.</p> | |
| <div class="pill"><span class="dot"></span> Local browser processing <span style="opacity:.45">·</span> ONNX Runtime</div> | |
| </header> | |
| <section class="panel"> | |
| <div class="section-head"><div class="step">01</div><div><h2>Source & settings</h2><div class="muted">Choose a recording and describe what you want to isolate.</div></div></div> | |
| <div class="file-wrap"><label for="sam-audio-file">Audio file</label><input id="sam-audio-file" type="file" accept=".wav,.mp3,.flac,.ogg,.m4a,audio/wav,audio/mpeg,audio/flac,audio/ogg,audio/mp4"><div class="muted" style="margin-top:8px">WAV, MP3, FLAC, OGG, or M4A · Maximum duration: 2 minutes</div></div> | |
| <div class="controls"> | |
| <div class="field prompt-field"><label for="sam-prompt">What should be isolated?</label><input id="sam-prompt" type="text" value="a person speaking" placeholder="e.g. a person speaking"></div> | |
| <div class="field"><label for="sam-provider">Processing device</label><select id="sam-provider"><option value="auto" selected>Automatic</option><option value="webgpu">GPU · WebGPU</option><option value="wasm">CPU · WASM</option></select></div> | |
| <div class="field"><label for="sam-steps">ODE steps</label><input id="sam-steps" type="number" min="1" max="16" value="8"></div> | |
| </div> | |
| <div class="muted" style="margin-top:9px">Automatic mode uses CPU on phones/tablets and WebGPU on desktop when available. More ODE steps can improve quality but take longer.</div> | |
| <div class="actions"><button id="sam-run">▶ Separate audio</button><button id="sam-clear-cache" class="secondary" type="button">Clear model cache</button><div id="sam-status-box" class="status" data-state="idle"><span class="status-dot"></span><span id="sam-status" aria-live="polite">Ready when you are</span></div></div> | |
| <div class="privacy"><svg viewBox="0 0 24 24" width="17" height="17" fill="none" stroke="currentColor" stroke-width="2"><rect x="4" y="10" width="16" height="11" rx="2"/><path d="M8 10V7a4 4 0 0 1 8 0v3"/></svg><div><strong>Your audio stays on your device.</strong> Inference runs locally in this browser. The Space hosts the interface and model downloads; your selected audio is not uploaded for processing.</div></div> | |
| </section> | |
| <section class="panel"> | |
| <div class="section-head"><div class="step">02</div><div><h2>Listen & compare</h2><div class="muted">Use the waveform to seek through each track. Download outputs as WAV files.</div></div></div> | |
| <div class="audio-grid"> | |
| <article class="audio-card"><div class="audio-card-head"><span class="audio-title">Original audio</span><span class="tag neutral">Input</span></div><div class="wave-wrap"><canvas id="sam-input-wave" aria-label="Original audio waveform"></canvas><div class="wave-empty" id="sam-input-empty">Choose an audio file to preview its waveform</div></div><div class="audio-meta"><span id="sam-input-duration">No audio loaded</span><span id="sam-input-format">Source</span></div><audio id="sam-input-player" controls preload="metadata"></audio><a class="download" id="sam-input-download" download hidden>↓ Download original</a></article> | |
| <article class="audio-card"><div class="audio-card-head"><span class="audio-title">Target</span><span class="tag">Isolated</span></div><div class="wave-wrap"><canvas id="sam-target-wave" aria-label="Target audio waveform"></canvas><div class="wave-empty" id="sam-target-empty">Run separation to create the target</div></div><div class="audio-meta"><span id="sam-target-duration">Waiting for result</span><span>48 kHz · WAV</span></div><audio id="sam-target-player" controls preload="metadata"></audio><a class="download" id="sam-target-download" download="target.wav" hidden>↓ Download target.wav</a></article> | |
| <article class="audio-card"><div class="audio-card-head"><span class="audio-title">Residual</span><span class="tag neutral">Everything else</span></div><div class="wave-wrap"><canvas id="sam-residual-wave" aria-label="Residual audio waveform"></canvas><div class="wave-empty" id="sam-residual-empty">Run separation to create the residual</div></div><div class="audio-meta"><span id="sam-residual-duration">Waiting for result</span><span>48 kHz · WAV</span></div><audio id="sam-residual-player" controls preload="metadata"></audio><a class="download" id="sam-residual-download" download="residual.wav" hidden>↓ Download residual.wav</a></article> | |
| </div> | |
| </section> | |
| <p class="footer-note">First use downloads model files and may take several minutes on CPU.</p> | |
| <div class="credits">Credits: <a href="https://github.com/facebookresearch/sam-audio" target="_blank" rel="noopener noreferrer">Meta's SAM Audio</a> <span aria-hidden="true">·</span> ONNX-converted models by <a href="https://huggingface.co/matbee/sam-audio-small-onnx" target="_blank" rel="noopener noreferrer">matbee</a></div> | |
| </div> | |
| """ | |
| CLIENT_JS = r""" | |
| async () => { | |
| const root = document.getElementById("sam-browser-app"); | |
| if (!root || root.dataset.initialized === "yes") return; | |
| root.dataset.initialized = "yes"; | |
| const $ = (id) => root.querySelector("#" + id); | |
| const MODEL_BASE = "https://huggingface.co/matbee/sam-audio-small-onnx/resolve/main"; | |
| const CACHE_NAME = "sam-audio-webgpu-models-v1"; | |
| const MAX_AUDIO_SECONDS = 120; | |
| const MODEL_FILES = { | |
| encoder: "dacvae_encoder.onnx", | |
| decoder: "dacvae_decoder.onnx", | |
| t5: "t5_encoder.onnx", | |
| dit: "dit_single_step.onnx" | |
| }; | |
| let ort = null; | |
| let tokenizer = null; | |
| let sessions = null; | |
| let sessionsProvider = null; | |
| let runtimeProvider = null; | |
| let lastAttemptProvider = null; | |
| let previousUrls = []; | |
| function setLog(message) { | |
| // Compact user-facing progress message; technical diagnostics are intentionally hidden. | |
| if (message) setStatus(String(message).split("\n")[0]); | |
| } | |
| function setStatus(message, state = "busy") { | |
| $("sam-status").textContent = message; | |
| $("sam-status-box").dataset.state = state; | |
| } | |
| const waveforms = {}; | |
| const colors = {input: "#8b879b", target: "#8b7cf6", residual: "#40b995"}; | |
| function formatDuration(seconds) { | |
| if (!Number.isFinite(seconds) || seconds < 0) return "—"; | |
| const mins = Math.floor(seconds / 60); | |
| const secs = Math.floor(seconds % 60); | |
| return mins ? `${mins}:${String(secs).padStart(2, "0")}` : `0:${String(secs).padStart(2, "0")}`; | |
| } | |
| function drawWaveform(canvasId, samples, kind) { | |
| const canvas = $(canvasId); | |
| if (!canvas) return; | |
| const rect = canvas.getBoundingClientRect(); | |
| if (!rect.width || !rect.height) return; | |
| const dpr = Math.max(1, Math.min(window.devicePixelRatio || 1, 2)); | |
| canvas.width = Math.round(rect.width * dpr); | |
| canvas.height = Math.round(rect.height * dpr); | |
| const ctx = canvas.getContext("2d"); | |
| ctx.setTransform(dpr, 0, 0, dpr, 0, 0); | |
| const width = rect.width, height = rect.height, mid = height / 2; | |
| ctx.clearRect(0, 0, width, height); | |
| ctx.strokeStyle = "rgba(128,128,150,.13)"; ctx.lineWidth = 1; | |
| ctx.beginPath(); ctx.moveTo(0, mid); ctx.lineTo(width, mid); ctx.stroke(); | |
| if (!samples || !samples.length) return; | |
| const bars = Math.max(1, Math.floor(width / 3)); | |
| const stride = Math.max(1, Math.floor(samples.length / bars)); | |
| const color = colors[kind] || colors.input; | |
| ctx.strokeStyle = color; ctx.lineWidth = Math.max(1, Math.min(2, width / bars * .7)); ctx.lineCap = "round"; | |
| ctx.beginPath(); | |
| for (let x = 0; x < bars; x++) { | |
| const start = x * stride; | |
| const end = Math.min(samples.length, start + stride); | |
| let peak = 0; | |
| for (let i = start; i < end; i++) peak = Math.max(peak, Math.abs(samples[i] || 0)); | |
| const amplitude = Math.max(1.25, peak * (height * .43)); | |
| const px = (x + .5) * width / bars; | |
| ctx.moveTo(px, mid - amplitude); ctx.lineTo(px, mid + amplitude); | |
| } | |
| ctx.stroke(); | |
| waveforms[canvasId] = {samples, kind}; | |
| updateWaveProgress(canvasId); | |
| } | |
| function updateWaveProgress(canvasId) { | |
| const canvas = $(canvasId); | |
| const player = $(canvasId.replace("-wave", "-player")); | |
| const entry = waveforms[canvasId]; | |
| if (!canvas || !entry || !player || !player.duration) return; | |
| const rect = canvas.getBoundingClientRect(); | |
| const dpr = Math.max(1, Math.min(window.devicePixelRatio || 1, 2)); | |
| const ctx = canvas.getContext("2d"); | |
| const x = Math.max(0, Math.min(rect.width, player.currentTime / player.duration * rect.width)); | |
| ctx.save(); ctx.setTransform(dpr, 0, 0, dpr, 0, 0); | |
| ctx.fillStyle = "rgba(255,255,255,.13)"; ctx.fillRect(0, 0, x, rect.height); | |
| ctx.restore(); | |
| } | |
| function attachWavePlayer(canvasId, playerId) { | |
| const canvas = $(canvasId), player = $(playerId); | |
| if (!canvas || !player) return; | |
| player.addEventListener("timeupdate", () => drawWaveform(canvasId, waveforms[canvasId]?.samples, waveforms[canvasId]?.kind || "input")); | |
| canvas.addEventListener("click", event => { | |
| if (!Number.isFinite(player.duration) || !player.duration) return; | |
| const rect = canvas.getBoundingClientRect(); | |
| player.currentTime = Math.max(0, Math.min(1, (event.clientX - rect.left) / rect.width)) * player.duration; | |
| }); | |
| } | |
| function setWave(id, samples, kind, durationId, emptyId) { | |
| const empty = $(emptyId); | |
| if (samples && samples.length) { | |
| if (empty) empty.hidden = true; | |
| drawWaveform(id, samples, kind); | |
| const duration = $(durationId); | |
| if (duration) duration.textContent = formatDuration(samples.length / 48000); | |
| } else { | |
| if (empty) empty.hidden = false; | |
| const canvas = $(id), ctx = canvas.getContext("2d"); | |
| ctx.clearRect(0, 0, canvas.width, canvas.height); | |
| delete waveforms[id]; | |
| } | |
| } | |
| function isMobileDevice() { | |
| // Prefer CPU on phones/tablets: mobile WebGPU drivers may expose the API | |
| // but fail during large model inference. Desktop users can still override manually. | |
| return /Android|iPhone|iPad|iPod|Mobile|Tablet/i.test(navigator.userAgent) || | |
| (navigator.userAgentData?.mobile === true); | |
| } | |
| async function chooseProvider() { | |
| const requested = $("sam-provider").value; | |
| if (requested === "webgpu" || requested === "wasm") return requested; | |
| if (isMobileDevice()) { | |
| setLog("Automatic device selection: mobile device detected. Using CPU (WASM) for compatibility."); | |
| return "wasm"; | |
| } | |
| try { | |
| if (navigator.gpu) { | |
| const adapter = await navigator.gpu.requestAdapter(); | |
| if (adapter) { | |
| setLog("Automatic device selection: desktop WebGPU adapter found. Using GPU (WebGPU)."); | |
| return "webgpu"; | |
| } | |
| } | |
| } catch (error) { | |
| setLog("WebGPU adapter check failed. Using CPU (WASM).\n" + String(error)); | |
| } | |
| setLog("Automatic device selection: WebGPU is unavailable. Using CPU (WASM)."); | |
| return "wasm"; | |
| } | |
| async function loadOrt() { | |
| if (!ort) { | |
| ort = await import("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.23.2/+esm"); | |
| if (ort.env?.wasm) { | |
| // Let ONNX Runtime choose its supported thread count. Multi-threading | |
| // is enabled only when the browser provides the required isolation. | |
| ort.env.wasm.numThreads = 0; | |
| ort.env.wasm.proxy = false; | |
| } | |
| if (ort.env) ort.env.logLevel = "warning"; | |
| } | |
| return ort; | |
| } | |
| async function loadTokenizer() { | |
| if (!tokenizer) { | |
| const transformers = await import("https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.2"); | |
| // Let Transformers.js persist tokenizer assets in its browser cache when supported. | |
| transformers.env.useBrowserCache = true; | |
| // Use the Transformers.js-compatible T5 tokenizer repository. | |
| tokenizer = await transformers.AutoTokenizer.from_pretrained("Xenova/t5-base"); | |
| } | |
| return tokenizer; | |
| } | |
| async function fetchCachedBytes(url, label) { | |
| // Cache Storage persists across page reloads. If unavailable or quota is exceeded, | |
| // continue with a normal network download rather than failing inference. | |
| let cache = null; | |
| if ("caches" in window) { | |
| try { cache = await caches.open(CACHE_NAME); } catch (_) {} | |
| } | |
| if (cache) { | |
| try { | |
| const cached = await cache.match(url); | |
| if (cached) { | |
| setLog("Using browser cache: " + label + "\nNo network download needed for this file."); | |
| return new Uint8Array(await cached.arrayBuffer()); | |
| } | |
| } catch (_) {} | |
| } | |
| setLog("Downloading " + label + "…\nThe file will be cached for future runs if browser storage permits."); | |
| const response = await fetch(url); | |
| if (!response.ok) throw new Error("Download failed (" + response.status + ") for " + url); | |
| if (cache) { | |
| try { | |
| // Cache a clone so the original response remains available for the ONNX runtime. | |
| await cache.put(url, response.clone()); | |
| } catch (error) { | |
| setLog("Downloaded " + label + ", but the browser could not persist it in cache (quota/private mode). Continuing without persistent cache.\n" + String(error)); | |
| } | |
| } | |
| return new Uint8Array(await response.arrayBuffer()); | |
| } | |
| async function configureRuntime(provider) { | |
| const runtime = await loadOrt(); | |
| if (runtimeProvider !== null && runtimeProvider !== provider) { | |
| throw new Error( | |
| "The CPU worker and GPU modes require different ONNX Runtime initialization. " + | |
| "Reload this page before switching between CPU and GPU, then select the desired mode before the first inference." | |
| ); | |
| } | |
| if (runtimeProvider === null) { | |
| if (runtime.env?.wasm) { | |
| // In CPU mode, run WASM inference in a proxy Web Worker so heavy | |
| // inference does not block the browser's UI thread. | |
| // WebGPU and the WASM proxy worker cannot be used together. | |
| runtime.env.wasm.proxy = provider === "wasm"; | |
| runtime.env.wasm.numThreads = 0; // automatic, subject to browser support | |
| } | |
| runtimeProvider = provider; | |
| } | |
| return runtime; | |
| } | |
| async function createSession(filename, provider) { | |
| const url = MODEL_BASE + "/" + filename; | |
| const runtime = await configureRuntime(provider); | |
| const modelBytes = await fetchCachedBytes(url, filename); | |
| const dataUrl = url + ".data"; | |
| const externalBytes = await fetchCachedBytes(dataUrl, filename + ".data (external weights)"); | |
| return runtime.InferenceSession.create(modelBytes, { | |
| executionProviders: [provider], | |
| graphOptimizationLevel: "all", | |
| externalData: [{ path: filename + ".data", data: externalBytes.buffer }] | |
| }); | |
| } | |
| async function clearModelCache() { | |
| if (sessions) { | |
| for (const session of Object.values(sessions)) { | |
| try { await session.release(); } catch (_) {} | |
| } | |
| sessions = null; | |
| sessionsProvider = null; | |
| } | |
| if ("caches" in window) { | |
| try { await caches.delete(CACHE_NAME); } catch (_) {} | |
| } | |
| setStatus("Model cache cleared · reload before next run", "success"); | |
| } | |
| async function loadModels(provider) { | |
| if (sessions && sessionsProvider === provider) return sessions; | |
| if (sessions) { | |
| for (const session of Object.values(sessions)) { | |
| try { await session.release(); } catch (_) {} | |
| } | |
| sessions = null; | |
| sessionsProvider = null; | |
| } | |
| const loaded = {}; | |
| const order = [ | |
| ["encoder", "dacvae_encoder.onnx"], | |
| ["t5", "t5_encoder.onnx"], | |
| ["dit", "dit_single_step.onnx"], | |
| ["decoder", "dacvae_decoder.onnx"] | |
| ]; | |
| try { | |
| for (const [key, filename] of order) { | |
| setStatus("Loading " + filename + "…"); | |
| setLog("Loading model component: " + filename + "\nThe browser may download a large external .onnx.data file. Keep this tab open."); | |
| loaded[key] = await createSession(filename, provider); | |
| setLog("Loaded " + filename + "\nInputs: " + loaded[key].inputNames.join(", ") + "\nOutputs: " + loaded[key].outputNames.join(", ")); | |
| } | |
| await loadTokenizer(); | |
| sessions = loaded; | |
| sessionsProvider = provider; | |
| return sessions; | |
| } catch (error) { | |
| for (const session of Object.values(loaded)) { | |
| try { await session.release(); } catch (_) {} | |
| } | |
| throw error; | |
| } | |
| } | |
| function getTensorMetadata(session, name) { | |
| try { return session.inputMetadata?.[name] || {}; } catch (_) { return {}; } | |
| } | |
| function floatTensor(data, dims) { | |
| // IMPORTANT: with ort.env.wasm.proxy=true, ONNX Runtime transfers input | |
| // ArrayBuffers to its worker. Passing the original typed array detaches it, | |
| // leaving length 0 for later ODE steps. Always pass a fresh copy. | |
| const copiedData = data instanceof Float32Array ? new Float32Array(data) : Float32Array.from(data); | |
| return new ort.Tensor("float32", copiedData, dims); | |
| } | |
| function int64Tensor(data, dims) { | |
| const values = Array.from(data, value => BigInt(value)); | |
| return new ort.Tensor("int64", BigInt64Array.from(values), dims); | |
| } | |
| function boolTensor(data, dims) { | |
| return new ort.Tensor("bool", Uint8Array.from(data, value => value ? 1 : 0), dims); | |
| } | |
| async function decodeAudio(file) { | |
| const bytes = await file.arrayBuffer(); | |
| const AudioContextClass = window.AudioContext || window.webkitAudioContext; | |
| const context = new AudioContextClass(); | |
| try { | |
| const decoded = await context.decodeAudioData(bytes.slice(0)); | |
| if (decoded.duration > MAX_AUDIO_SECONDS) { | |
| throw new Error("Audio must be 2 minutes or shorter. Selected file is " + formatDuration(decoded.duration) + "."); | |
| } | |
| const offline = new OfflineAudioContext(1, Math.ceil(decoded.duration * 48000), 48000); | |
| const source = offline.createBufferSource(); | |
| const mono = offline.createBuffer(1, decoded.length, decoded.sampleRate); | |
| const channels = []; | |
| for (let i = 0; i < decoded.numberOfChannels; i++) channels.push(decoded.getChannelData(i)); | |
| const monoData = mono.getChannelData(0); | |
| for (let i = 0; i < monoData.length; i++) { | |
| let sum = 0; | |
| for (const channel of channels) sum += channel[i]; | |
| monoData[i] = sum / channels.length; | |
| } | |
| source.buffer = mono; | |
| source.connect(offline.destination); | |
| source.start(); | |
| const resampled = await offline.startRendering(); | |
| return {samples: new Float32Array(resampled.getChannelData(0)), sampleRate: 48000}; | |
| } finally { | |
| await context.close().catch(() => {}); | |
| } | |
| } | |
| async function runSession(session, outputName, feeds) { | |
| const outputs = await session.run(feeds); | |
| const tensor = outputs[outputName] || outputs[session.outputNames[0]]; | |
| if (!tensor) throw new Error("Expected output not found: " + outputName); | |
| return tensor.data; | |
| } | |
| async function encodeAudio(samples) { | |
| const session = sessions.encoder; | |
| // In WASM proxy mode, ONNX Runtime may transfer/detach the tensor's | |
| // ArrayBuffer when sending it to the worker. Never give it audio.samples | |
| // directly: keep the decoded source intact for duration checks and output trimming. | |
| const inputSamples = new Float32Array(samples); | |
| const input = new ort.Tensor("float32", inputSamples, [1, 1, inputSamples.length]); | |
| const out = await session.run({audio: input}); | |
| const tensor = out.latent_features || out[session.outputNames[0]]; | |
| return {data: new Float32Array(tensor.data), dims: tensor.dims}; | |
| } | |
| async function encodeText(text) { | |
| const tok = await loadTokenizer(); | |
| const encoded = await tok(text, { | |
| padding: "max_length", | |
| truncation: true, | |
| max_length: 77 | |
| }); | |
| const idsRaw = encoded.input_ids?.data ?? encoded.input_ids; | |
| const maskRaw = encoded.attention_mask?.data ?? encoded.attention_mask; | |
| const ids = Array.from(idsRaw); | |
| const mask = Array.from(maskRaw); | |
| const idsTensor = int64Tensor(ids, [1, ids.length]); | |
| const maskTensor = int64Tensor(mask, [1, mask.length]); | |
| const out = await sessions.t5.run({input_ids: idsTensor, attention_mask: maskTensor}); | |
| const tensor = out.hidden_states || out[sessions.t5.outputNames[0]]; | |
| return { | |
| features: new Float32Array(tensor.data), | |
| featureDims: tensor.dims, | |
| mask: Uint8Array.from(mask, value => value ? 1 : 0), | |
| maskDims: [1, mask.length] | |
| }; | |
| } | |
| async function runDit(noisy, timeValue, audioFeatures, text, sequenceLength) { | |
| const batch = 1; | |
| const zerosInt = new BigInt64Array(sequenceLength); | |
| const mask = new Uint8Array(sequenceLength).fill(1); | |
| const video = new Float32Array(1024 * sequenceLength); | |
| const feeds = { | |
| noisy_audio: floatTensor(noisy, [batch, sequenceLength, 256]), | |
| time: floatTensor(new Float32Array([timeValue]), [1]), | |
| audio_features: floatTensor(audioFeatures, [batch, sequenceLength, 256]), | |
| text_features: floatTensor(text.features, text.featureDims), | |
| text_mask: boolTensor(text.mask, text.maskDims), | |
| masked_video_features: floatTensor(video, [batch, 1024, sequenceLength]), | |
| anchor_ids: new ort.Tensor("int64", zerosInt, [batch, sequenceLength]), | |
| anchor_alignment: new ort.Tensor("int64", new BigInt64Array(sequenceLength), [batch, sequenceLength]), | |
| audio_pad_mask: boolTensor(mask, [batch, sequenceLength]) | |
| }; | |
| const outputs = await sessions.dit.run(feeds); | |
| const tensor = outputs.velocity || outputs[sessions.dit.outputNames[0]]; | |
| return new Float32Array(tensor.data); | |
| } | |
| async function runOde(initial, audioFeatures, text, steps, sequenceLength) { | |
| let x = new Float32Array(initial); | |
| const dt = 1 / steps; | |
| for (let i = 0; i < steps; i++) { | |
| setStatus("Denoising step " + (i + 1) + "/" + steps + " (1/2)…"); | |
| const t = i * dt; | |
| const k1 = await runDit(x, t, audioFeatures, text, sequenceLength); | |
| const midpoint = new Float32Array(x.length); | |
| for (let j = 0; j < x.length; j++) midpoint[j] = x[j] + (dt / 2) * k1[j]; | |
| setStatus("Denoising step " + (i + 1) + "/" + steps + " (2/2)…"); | |
| const k2 = await runDit(midpoint, t + dt / 2, audioFeatures, text, sequenceLength); | |
| for (let j = 0; j < x.length; j++) x[j] += dt * k2[j]; | |
| // Yield so status/progress updates can paint between long CPU steps. | |
| await new Promise(resolve => setTimeout(resolve, 0)); | |
| } | |
| return x; | |
| } | |
| async function decodeLatent(latent, latentChannels, timeSteps) { | |
| const session = sessions.decoder; | |
| const chunkSize = 25; | |
| const hopLength = 1920; | |
| const chunks = []; | |
| for (let start = 0; start < timeSteps; start += chunkSize) { | |
| const actual = Math.min(chunkSize, timeSteps - start); | |
| const chunk = new Float32Array(latentChannels * chunkSize); | |
| for (let c = 0; c < latentChannels; c++) { | |
| for (let t = 0; t < actual; t++) { | |
| chunk[c * chunkSize + t] = latent[c * timeSteps + start + t]; | |
| } | |
| } | |
| const out = await session.run({ | |
| latent_features: floatTensor(chunk, [1, latentChannels, chunkSize]) | |
| }); | |
| const tensor = out.waveform || out[session.outputNames[0]]; | |
| const data = tensor.data; | |
| const count = actual * hopLength; | |
| chunks.push(new Float32Array(data.slice(0, count))); | |
| setStatus("Decoding audio " + Math.min(start + chunkSize, timeSteps) + "/" + timeSteps + " latent frames…"); | |
| } | |
| const totalLength = chunks.reduce((sum, part) => sum + part.length, 0); | |
| const output = new Float32Array(totalLength); | |
| let offset = 0; | |
| for (const chunk of chunks) { output.set(chunk, offset); offset += chunk.length; } | |
| return output; | |
| } | |
| function wavBlob(samples, sampleRate) { | |
| const buffer = new ArrayBuffer(44 + samples.length * 2); | |
| const view = new DataView(buffer); | |
| function writeString(offset, value) { for (let i = 0; i < value.length; i++) view.setUint8(offset + i, value.charCodeAt(i)); } | |
| writeString(0, "RIFF"); view.setUint32(4, 36 + samples.length * 2, true); | |
| writeString(8, "WAVE"); writeString(12, "fmt "); view.setUint32(16, 16, true); | |
| view.setUint16(20, 1, true); view.setUint16(22, 1, true); view.setUint32(24, sampleRate, true); | |
| view.setUint32(28, sampleRate * 2, true); view.setUint16(32, 2, true); view.setUint16(34, 16, true); | |
| writeString(36, "data"); view.setUint32(40, samples.length * 2, true); | |
| for (let i = 0; i < samples.length; i++) { | |
| const sample = Math.max(-1, Math.min(1, samples[i])); | |
| view.setInt16(44 + i * 2, sample < 0 ? sample * 32768 : sample * 32767, true); | |
| } | |
| return new Blob([buffer], {type:"audio/wav"}); | |
| } | |
| function showOutput(playerId, linkId, samples, sampleRate, filename, waveId, kind, durationId, emptyId) { | |
| const url = URL.createObjectURL(wavBlob(samples, sampleRate)); | |
| previousUrls.push(url); | |
| const player = $(playerId); | |
| player.src = url; | |
| player.load(); | |
| const link = $(linkId); link.href = url; link.download = filename; link.hidden = false; | |
| setWave(waveId, samples, kind, durationId, emptyId); | |
| } | |
| function showInput(file, samples) { | |
| const url = URL.createObjectURL(file); | |
| previousUrls.push(url); | |
| const player = $("sam-input-player"); player.src = url; player.load(); | |
| const link = $("sam-input-download"); link.href = url; link.download = file.name; link.hidden = false; link.textContent = "↓ Download original"; | |
| $("sam-input-format").textContent = file.name.length > 24 ? file.name.slice(0, 21) + "…" : file.name; | |
| setWave("sam-input-wave", samples, "input", "sam-input-duration", "sam-input-empty"); | |
| } | |
| async function separate(providerOverride = null) { | |
| const file = $("sam-audio-file").files?.[0]; | |
| const prompt = $("sam-prompt").value.trim(); | |
| const steps = Number.parseInt($("sam-steps").value, 10); | |
| const provider = providerOverride || await chooseProvider(); | |
| lastAttemptProvider = provider; | |
| const deviceLabel = provider === "wasm" ? "CPU (WASM)" : "GPU (WebGPU)"; | |
| if (!file) throw new Error("Choose an audio file first."); | |
| if (!prompt) throw new Error("Enter a text prompt."); | |
| if (!Number.isInteger(steps) || steps < 1 || steps > 16) throw new Error("ODE steps must be between 1 and 16."); | |
| const audio = await decodeAudio(file); | |
| const expectedSamples = audio.samples.length; | |
| if (expectedSamples > 48000 * MAX_AUDIO_SECONDS) throw new Error("Audio must be 2 minutes or shorter."); | |
| setStatus("Loading models for " + deviceLabel + "…"); | |
| if (provider === "wasm") { | |
| setLog("CPU mode selected. ONNX Runtime will use WebAssembly in a background worker where supported, keeping the page responsive. This may take longer than GPU inference."); | |
| } | |
| await loadModels(provider); | |
| setStatus("Encoding audio on " + deviceLabel + "…"); | |
| const encoded = await encodeAudio(audio.samples); | |
| if (encoded.dims.length !== 3) throw new Error("Unexpected encoder output dimensions: " + JSON.stringify(encoded.dims)); | |
| const [batch, latentChannels, timeSteps] = encoded.dims; | |
| if (batch !== 1 || latentChannels !== 128) { | |
| throw new Error("Unexpected audio latent shape " + JSON.stringify(encoded.dims) + "; expected [1, 128, time]."); | |
| } | |
| setStatus("Encoding text…"); | |
| const text = await encodeText(prompt); | |
| const mixture = new Float32Array(timeSteps * 256); | |
| for (let t = 0; t < timeSteps; t++) { | |
| for (let c = 0; c < latentChannels; c++) { | |
| const value = encoded.data[c * timeSteps + t]; | |
| mixture[t * 256 + c] = value; | |
| mixture[t * 256 + latentChannels + c] = value; | |
| } | |
| } | |
| const initial = new Float32Array(timeSteps * 256); | |
| // Standard normal noise, matching the reference implementation's randn initialization. | |
| for (let i = 0; i < initial.length; i += 2) { | |
| const u1 = Math.max(Number.EPSILON, Math.random()); | |
| const u2 = Math.random(); | |
| const radius = Math.sqrt(-2 * Math.log(u1)); | |
| initial[i] = radius * Math.cos(2 * Math.PI * u2); | |
| if (i + 1 < initial.length) initial[i + 1] = radius * Math.sin(2 * Math.PI * u2); | |
| } | |
| const result = await runOde(initial, mixture, text, steps, timeSteps); | |
| const targetLatent = new Float32Array(latentChannels * timeSteps); | |
| const residualLatent = new Float32Array(latentChannels * timeSteps); | |
| for (let t = 0; t < timeSteps; t++) { | |
| for (let c = 0; c < latentChannels; c++) { | |
| targetLatent[c * timeSteps + t] = result[t * 256 + c]; | |
| residualLatent[c * timeSteps + t] = result[t * 256 + latentChannels + c]; | |
| } | |
| } | |
| setStatus("Decoding target…"); | |
| const target = await decodeLatent(targetLatent, latentChannels, timeSteps); | |
| setStatus("Decoding residual…"); | |
| const residual = await decodeLatent(residualLatent, latentChannels, timeSteps); | |
| if (audio.samples.length !== expectedSamples) { | |
| throw new Error("Internal audio buffer was unexpectedly detached or resized: expected " + expectedSamples + " samples, got " + audio.samples.length + "."); | |
| } | |
| showOutput("sam-target-player", "sam-target-download", target.slice(0, expectedSamples), 48000, "target.wav", "sam-target-wave", "target", "sam-target-duration", "sam-target-empty"); | |
| showOutput("sam-residual-player", "sam-residual-download", residual.slice(0, expectedSamples), 48000, "residual.wav", "sam-residual-wave", "residual", "sam-residual-duration", "sam-residual-empty"); | |
| setStatus("Separation complete · " + deviceLabel, "success"); | |
| } | |
| attachWavePlayer("sam-input-wave", "sam-input-player"); | |
| attachWavePlayer("sam-target-wave", "sam-target-player"); | |
| attachWavePlayer("sam-residual-wave", "sam-residual-player"); | |
| $("sam-audio-file").addEventListener("change", async () => { | |
| const file = $("sam-audio-file").files?.[0]; | |
| if (!file) return; | |
| for (const url of previousUrls) URL.revokeObjectURL(url); | |
| previousUrls = []; | |
| for (const [playerId, linkId, waveId, durationId, emptyId, waitingText] of [ | |
| ["sam-input-player", "sam-input-download", "sam-input-wave", "sam-input-duration", "sam-input-empty", "No audio loaded"], | |
| ["sam-target-player", "sam-target-download", "sam-target-wave", "sam-target-duration", "sam-target-empty", "Waiting for result"], | |
| ["sam-residual-player", "sam-residual-download", "sam-residual-wave", "sam-residual-duration", "sam-residual-empty", "Waiting for result"] | |
| ]) { | |
| const player = $(playerId); player.pause(); player.removeAttribute("src"); player.load(); | |
| const link = $(linkId); link.removeAttribute("href"); link.hidden = true; | |
| $(durationId).textContent = waitingText; | |
| $(emptyId).hidden = false; | |
| const canvas = $(waveId); const ctx = canvas.getContext("2d"); ctx.clearRect(0, 0, canvas.width, canvas.height); | |
| delete waveforms[waveId]; | |
| } | |
| try { | |
| setStatus("Preparing audio preview…", "busy"); | |
| const preview = await decodeAudio(file); | |
| showInput(file, preview.samples); | |
| setStatus("Source ready · " + formatDuration(preview.samples.length / preview.sampleRate), "idle"); | |
| } catch (error) { | |
| setStatus("Could not preview this audio: " + String(error?.message || error), "error"); | |
| } | |
| }); | |
| window.addEventListener("resize", () => { | |
| for (const [id, entry] of Object.entries(waveforms)) drawWaveform(id, entry.samples, entry.kind); | |
| }); | |
| $("sam-clear-cache").addEventListener("click", async () => { | |
| const button = $("sam-clear-cache"); | |
| button.disabled = true; | |
| try { await clearModelCache(); } | |
| finally { button.disabled = false; } | |
| }); | |
| $("sam-run").addEventListener("click", async () => { | |
| const button = $("sam-run"); | |
| button.disabled = true; | |
| try { | |
| await separate(); | |
| } catch (error) { | |
| const message = String(error?.message || error); | |
| const isWebGpuBufferFailure = lastAttemptProvider === "webgpu" && | |
| /createBuffer failed.*mappedAtCreation|mappedAtCreation.*too large/i.test(message); | |
| if (isWebGpuBufferFailure) { | |
| const retryOnCpu = window.confirm( | |
| "WebGPU failed to allocate a GPU buffer on this device. This can be a browser/driver or GPU-memory limitation.\n\nRetry the same audio using CPU (WASM)? This may be slower." | |
| ); | |
| if (retryOnCpu) { | |
| $("sam-provider").value = "wasm"; | |
| setStatus("Retrying on CPU…"); | |
| setLog("WebGPU buffer allocation failed. Retrying locally with CPU (WASM).\nThe audio is not uploaded to the server."); | |
| try { | |
| await separate("wasm"); | |
| setStatus("Separation complete · CPU fallback used", "success"); | |
| } catch (cpuError) { | |
| setStatus("CPU retry failed: " + String(cpuError?.message || cpuError), "error"); | |
| } | |
| } else { | |
| setStatus("GPU allocation failed · try CPU (WASM)", "error"); | |
| } | |
| } else { | |
| setStatus("Separation failed: " + message, "error"); | |
| } | |
| } finally { | |
| button.disabled = false; | |
| } | |
| }); | |
| } | |
| """ | |
| def zerogpu_startup_marker(): | |
| return "No server inference; browser WebGPU prototype." | |
| with gr.Blocks(title="SAM Audio Studio", theme=gr.themes.Soft(primary_hue="violet", secondary_hue="slate"), css="footer {visibility: hidden}") as demo: | |
| gr.HTML(PAGE) | |
| demo.load(fn=None, inputs=None, outputs=None, js=CLIENT_JS) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) |