Update app.py

app.py

@@ -19,8 +19,17 @@ PAGE = r"""
 #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],#sam-browser-app input[type=text],#sam-browser-app input[type=number],#sam-browser-app select,#sam-browser-app textarea{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 textarea{width:100%;min-height:86px;resize:vertical;padding:11px 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:13px;line-height:1.5}
+#sam-browser-app .advanced{margin-top:18px;border:1px solid var(--border-color-primary);border-radius:13px;overflow:hidden}
+#sam-browser-app .advanced summary{cursor:pointer;padding:13px 15px;font-size:13px;font-weight:750;background:var(--background-fill-secondary)}
+#sam-browser-app .advanced-body{padding:15px;display:grid;gap:15px}
+#sam-browser-app .option-row{display:flex;gap:10px;align-items:flex-start}
+#sam-browser-app .option-row input{margin-top:3px;accent-color:var(--sam-accent)}
+#sam-browser-app .option-row label{margin:0;font-weight:750}
+#sam-browser-app .advanced-grid{display:grid;grid-template-columns:minmax(0,1fr) 120px;gap:12px;align-items:end}
+#sam-browser-app .threshold{width:100%;accent-color:var(--sam-accent)}
 #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}
@@ -53,7 +62,7 @@ PAGE = r"""
 #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 28px;line-height:1.5}
 @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%}}
+@media(max-width:480px){#sam-browser-app .advanced-grid{grid-template-columns:1fr}#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>
@@ -63,13 +72,22 @@ PAGE = r"""
 </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="audio/*,.wav,.mp3,.flac,.ogg,.m4a"><div class="muted" style="margin-top:8px">WAV, MP3, FLAC, OGG, or M4A · Recommended: a short clip for the first run</div></div>
+  <div class="file-wrap"><label for="sam-audio-file">Audio file</label><input id="sam-audio-file" type="file" accept="audio/*,.wav,.mp3,.flac,.ogg,.m4a"><div class="muted" style="margin-top:8px">WAV, MP3, FLAC, OGG, or M4A · Maximum duration: 60 seconds</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="2"></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>
+  <details class="advanced">
+    <summary>Extra settings · Span prediction & anchors</summary>
+    <div class="advanced-body">
+      <div class="option-row"><input id="sam-predict-spans" type="checkbox"><div><label for="sam-predict-spans">Automatic Span Prediction</label><div class="muted">Detect when the described sound occurs and guide separation toward those time ranges. First use may download about 5.8 GB of additional PEAFrame weights plus tokenizer assets; browser/device memory limits may prevent loading them.</div></div></div>
+      <div class="advanced-grid" id="sam-span-threshold-wrap" hidden><div><label for="sam-span-threshold">Detection threshold</label><input id="sam-span-threshold" class="threshold" type="range" min="0.05" max="0.95" step="0.05" value="0.30"></div><div id="sam-span-threshold-value" class="muted">0.30</div></div>
+      <div class="option-row"><input id="sam-manual-anchors-enabled" type="checkbox"><div><label for="sam-manual-anchors-enabled">Manual Anchors</label><div class="muted">Define positive ranges to focus on (+) or negative ranges to ignore (−). Each line uses: sign start_seconds end_seconds.</div></div></div>
+      <div id="sam-manual-anchors-wrap" hidden><label for="sam-manual-anchors">Anchor ranges</label><textarea id="sam-manual-anchors" spellcheck="false" placeholder="+ 4.5 7.0&#10;+ 12.0 15.5&#10;- 0.0 3.0"></textarea><div class="muted">Example: <code>+ 4.5 7.0</code> focuses on 4.5–7 seconds. Manual anchors take priority over automatic prediction.</div></div>
+    </div>
+  </details>
   <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>
@@ -101,6 +119,9 @@ async () => {
   };
   let ort = null;
   let tokenizer = null;
+  let transformersModule = null;
+  let spanTokenizer = null;
+  let spanConfig = null;
   let sessions = null;
   let sessionsProvider = null;
   let runtimeProvider = null;
@@ -235,14 +256,160 @@ async () => {
   }
   async function loadTokenizer() {
     if (!tokenizer) {
-      const transformers = await import("https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.2");
+      transformersModule = 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;
+      transformersModule.env.useBrowserCache = true;
       // Use the Transformers.js-compatible T5 tokenizer repository.
-      tokenizer = await transformers.AutoTokenizer.from_pretrained("Xenova/t5-base");
+      tokenizer = await transformersModule.AutoTokenizer.from_pretrained("Xenova/t5-base");
     }
     return tokenizer;
   }
+  async function loadSpanTokenizer() {
+    if (!transformersModule) {
+      transformersModule = await import("https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.2");
+      transformersModule.env.useBrowserCache = true;
+    }
+    if (!spanTokenizer) {
+      // Load the exact tokenizer files from the model repo through our shared
+      // Cache Storage, then instantiate the tokenizer directly from its JSON.
+      const tokenizerBase = MODEL_BASE + "/peaframe_tokenizer";
+      const tokenizerJsonBytes = await fetchCachedBytes(tokenizerBase + "/tokenizer.json", "PEAFrame tokenizer · tokenizer.json");
+      const tokenizerConfigBytes = await fetchCachedBytes(tokenizerBase + "/tokenizer_config.json", "PEAFrame tokenizer · tokenizer_config.json");
+      await fetchCachedBytes(tokenizerBase + "/special_tokens_map.json", "PEAFrame tokenizer · special_tokens_map.json");
+      const tokenizerJson = JSON.parse(new TextDecoder().decode(tokenizerJsonBytes));
+      const tokenizerConfig = JSON.parse(new TextDecoder().decode(tokenizerConfigBytes));
+      if (typeof transformersModule.PreTrainedTokenizer !== "function") {
+        throw new Error("This Transformers.js build does not expose PreTrainedTokenizer.");
+      }
+      spanTokenizer = new transformersModule.PreTrainedTokenizer(tokenizerJson, tokenizerConfig);
+    }
+    return spanTokenizer;
+  }
+  async function loadSpanPredictor(provider) {
+    if (!sessions) throw new Error("Load the base audio models before span prediction.");
+    if (sessions.peaframe) return sessions.peaframe;
+    setStatus("Loading optional span-prediction model…");
+    setLog("Automatic Span Prediction requires the optional PEAFrame ONNX model and external weights. First download can be several GB; assets are cached in this browser.");
+    const configBytes = await fetchCachedBytes(MODEL_BASE + "/peaframe_config.json", "PEAFrame configuration");
+    spanConfig = JSON.parse(new TextDecoder().decode(configBytes));
+    const session = await createSession("peaframe.onnx", provider);
+    try {
+      await loadSpanTokenizer();
+    } catch (error) {
+      try { await session.release(); } catch (_) {}
+      throw new Error("Could not load the PEAFrame tokenizer in this browser. Manual Anchors remain available. Details: " + String(error?.message || error));
+    }
+    sessions.peaframe = session;
+    return session;
+  }
+  function sigmoid(value) {
+    if (value >= 0) { const z = Math.exp(-value); return 1 / (1 + z); }
+    const z = Math.exp(value); return z / (1 + z);
+  }
+  async function predictSpans(samples, prompt, threshold, provider) {
+    const session = await loadSpanPredictor(provider);
+    const tok = await loadSpanTokenizer();
+    const encoded = await tok(prompt, {padding: true, truncation: true, max_length: 512});
+    const rawIds = encoded.input_ids?.data ?? encoded.input_ids;
+    const rawMask = encoded.attention_mask?.data ?? encoded.attention_mask;
+    const ids = Array.from(rawIds, Number);
+    const attention = Array.from(rawMask, Number);
+    const expectedSamples = 160000; // fixed PEAFrame ONNX input: ~3.33 s at 48 kHz
+    const hopLength = Number(spanConfig?.hop_length || 1920);
+    const sampleRate = Number(spanConfig?.sampling_rate || 48000);
+    const stride = expectedSamples >> 1;
+    const probabilityChunks = [];
+    const starts = [];
+    const startsToProcess = samples.length <= expectedSamples ? [0] : Array.from({length: Math.ceil((samples.length - expectedSamples) / stride) + 1}, (_, i) => i * stride);
+    for (let ci = 0; ci < startsToProcess.length; ci++) {
+      const start = startsToProcess[ci];
+      const end = Math.min(start + expectedSamples, samples.length);
+      const chunk = new Float32Array(expectedSamples);
+      chunk.set(samples.subarray(start, end));
+      const output = await session.run({
+        input_ids: int64Tensor(ids, [1, ids.length]),
+        input_values: floatTensor(chunk, [1, 1, expectedSamples]),
+        attention_mask: int64Tensor(attention, [1, attention.length])
+      });
+      const audioEmbeds = output.audio_embeds || output[session.outputNames[0]];
+      const textEmbeds = output.text_embeds || output[session.outputNames[1]];
+      if (!audioEmbeds || !textEmbeds) throw new Error("PEAFrame returned unexpected outputs; expected audio_embeds and text_embeds.");
+      const frames = audioEmbeds.dims[1];
+      const dim = audioEmbeds.dims[2];
+      const ae = audioEmbeds.data, te = textEmbeds.data;
+      const probabilities = new Float32Array(frames);
+      for (let frame = 0; frame < frames; frame++) {
+        let dot = 0;
+        const offset = frame * dim;
+        for (let d = 0; d < dim; d++) dot += ae[offset + d] * te[d];
+        probabilities[frame] = sigmoid(dot * Number(spanConfig?.logit_scale ?? 2.298705816268921) + Number(spanConfig?.logit_bias ?? -10.002328872680664));
+      }
+      const usableFrames = Math.min(frames, Math.ceil((end - start) / hopLength));
+      probabilityChunks.push(probabilities.slice(0, usableFrames));
+      starts.push(start);
+      setStatus("Predicting sound spans " + (ci + 1) + "/" + startsToProcess.length + "…");
+      await new Promise(resolve => setTimeout(resolve, 0));
+    }
+    const totalFrames = Math.ceil(samples.length / hopLength);
+    const merged = new Float32Array(totalFrames);
+    const counts = new Uint16Array(totalFrames);
+    for (let ci = 0; ci < probabilityChunks.length; ci++) {
+      const chunkStart = Math.floor(starts[ci] / hopLength);
+      const probs = probabilityChunks[ci];
+      for (let j = 0; j < probs.length && chunkStart + j < totalFrames; j++) {
+        merged[chunkStart + j] += probs[j];
+        counts[chunkStart + j]++;
+      }
+    }
+    for (let i = 0; i < merged.length; i++) merged[i] = counts[i] ? merged[i] / counts[i] : 0;
+    const spans = [];
+    let active = false, startFrame = 0;
+    for (let i = 0; i < merged.length; i++) {
+      const detected = merged[i] > threshold;
+      if (detected && !active) { startFrame = i; active = true; }
+      else if (!detected && active) {
+        spans.push([startFrame * hopLength / sampleRate, i * hopLength / sampleRate]);
+        active = false;
+      }
+    }
+    if (active) spans.push([startFrame * hopLength / sampleRate, samples.length / sampleRate]);
+    return spans;
+  }
+  function parseManualAnchors(raw, duration) {
+    const lines = raw.split(/\r?\n/).map(line => line.trim()).filter(Boolean);
+    const spans = [];
+    for (let index = 0; index < lines.length; index++) {
+      const parts = lines[index].split(/[\s,;]+/);
+      if (parts.length !== 3 || !["+", "-"].includes(parts[0])) {
+        throw new Error("Invalid manual anchor on line " + (index + 1) + '. Use: + 4.5 7.0 or - 0 3.0.');
+      }
+      const start = Number(parts[1]), end = Number(parts[2]);
+      if (!Number.isFinite(start) || !Number.isFinite(end) || start < 0 || end <= start || end > duration) {
+        throw new Error("Invalid time range on manual anchor line " + (index + 1) + ". Times must be within the audio duration and end must be greater than start.");
+      }
+      spans.push({sign: parts[0], start, end});
+    }
+    if (!spans.length) throw new Error("Manual Anchors is enabled, but no valid anchor ranges were entered.");
+    return spans;
+  }
+  function buildAnchorTensors(spans, sequenceLength) {
+    const ids = [0, 3]; // <null>=0, <pad>=3, as in the Python reference
+    const alignment = new BigInt64Array(sequenceLength);
+    alignment.fill(1); // unanchored frames map to <pad>
+    for (const span of spans) {
+      const startFrame = Math.max(0, Math.min(sequenceLength, Math.floor(span.start * 48000 / 1920)));
+      const endFrame = Math.max(0, Math.min(sequenceLength, Math.floor(span.end * 48000 / 1920)));
+      if (startFrame >= endFrame) continue;
+      const anchorIndex = ids.length;
+      alignment.fill(BigInt(anchorIndex), startFrame, endFrame);
+      ids.push(span.sign === "-" ? 2 : 1);
+    }
+    return {
+      ids: new ort.Tensor("int64", BigInt64Array.from(ids, value => BigInt(value)), [1, ids.length]),
+      alignment: new ort.Tensor("int64", alignment, [1, sequenceLength])
+    };
+  }
+
   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.
@@ -312,6 +479,10 @@ async () => {
       sessions = null;
       sessionsProvider = null;
     }
+    spanTokenizer = null;
+    spanConfig = null;
+    tokenizer = null;
+    transformersModule = null;
     if ("caches" in window) {
       try { await caches.delete(CACHE_NAME); } catch (_) {}
     }
@@ -433,9 +604,10 @@ async () => {
       maskDims: [1, mask.length]
     };
   }
-  async function runDit(noisy, timeValue, audioFeatures, text, sequenceLength) {
+  async function runDit(noisy, timeValue, audioFeatures, text, sequenceLength, anchors = null) {
     const batch = 1;
-    const zerosInt = new BigInt64Array(sequenceLength);
+    const anchorIds = anchors?.ids || new ort.Tensor("int64", BigInt64Array.from([0n, 3n]), [1, 2]);
+    const anchorAlignment = anchors?.alignment || new ort.Tensor("int64", new BigInt64Array(sequenceLength), [1, sequenceLength]);
     const mask = new Uint8Array(sequenceLength).fill(1);
     const video = new Float32Array(1024 * sequenceLength);
     const feeds = {
@@ -445,25 +617,25 @@ async () => {
       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]),
+      anchor_ids: anchorIds,
+      anchor_alignment: anchorAlignment,
       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) {
+  async function runOde(initial, audioFeatures, text, steps, sequenceLength, anchors = null) {
     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 k1 = await runDit(x, t, audioFeatures, text, sequenceLength, anchors);
       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);
+      const k2 = await runDit(midpoint, t + dt / 2, audioFeatures, text, sequenceLength, anchors);
       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));
@@ -542,7 +714,11 @@ async () => {
     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 * 20) throw new Error("For this first prototype, use audio no longer than 20 seconds.");
+    if (expectedSamples > 48000 * 60) throw new Error("Audio must be 60 seconds or shorter.");
+    const predictSpansEnabled = $("sam-predict-spans").checked;
+    const manualAnchorsEnabled = $("sam-manual-anchors-enabled").checked;
+    const manualSpans = manualAnchorsEnabled ? parseManualAnchors($("sam-manual-anchors").value, expectedSamples / 48000) : [];
+    const spanThreshold = Number($("sam-span-threshold").value);
     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.");
@@ -574,7 +750,22 @@ async () => {
       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);
+    let anchorSpans = [];
+    if (manualAnchorsEnabled) {
+      anchorSpans = manualSpans;
+      setStatus("Applying " + anchorSpans.length + " manual anchor range(s)…");
+    } else if (predictSpansEnabled) {
+      setStatus("Detecting target sound spans…");
+      const detected = await predictSpans(audio.samples, prompt, spanThreshold, provider);
+      anchorSpans = detected.map(([start, end]) => ({sign: "+", start, end}));
+      if (anchorSpans.length) {
+        setStatus("Detected " + anchorSpans.length + " sound span(s). Preparing separation…");
+      } else {
+        setStatus("No matching spans detected · continuing without span guidance…");
+      }
+    }
+    const anchors = anchorSpans.length ? buildAnchorTensors(anchorSpans, timeSteps) : null;
+    const result = await runOde(initial, mixture, text, steps, timeSteps, anchors);
     const targetLatent = new Float32Array(latentChannels * timeSteps);
     const residualLatent = new Float32Array(latentChannels * timeSteps);
     for (let t = 0; t < timeSteps; t++) {
@@ -626,6 +817,15 @@ async () => {
   window.addEventListener("resize", () => {
     for (const [id, entry] of Object.entries(waveforms)) drawWaveform(id, entry.samples, entry.kind);
   });
+  $("sam-predict-spans").addEventListener("change", () => {
+    $("sam-span-threshold-wrap").hidden = !$("sam-predict-spans").checked;
+  });
+  $("sam-manual-anchors-enabled").addEventListener("change", () => {
+    $("sam-manual-anchors-wrap").hidden = !$("sam-manual-anchors-enabled").checked;
+  });
+  $("sam-span-threshold").addEventListener("input", () => {
+    $("sam-span-threshold-value").textContent = Number($("sam-span-threshold").value).toFixed(2);
+  });
   $("sam-clear-cache").addEventListener("click", async () => {
     const button = $("sam-clear-cache");
     button.disabled = true;