app.py

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">▶ &nbsp; 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;
    }
  });
}
"""
 
@spaces.GPU
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)