Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-
+ 
 PAGE = r"""
 <div id="sam-browser-app" style="font-family:system-ui,sans-serif;max-width:1050px;margin:auto">
 <style>
@@ -18,8 +18,8 @@ PAGE = r"""
 #sam-browser-app .out{min-width:0}
 @media(max-width:650px){#sam-browser-app .outputs{grid-template-columns:1fr}}
 </style>
-<h1>SAM Audio Small — Browser WebGPU inference prototype</h1>
-<p>This is a browser-side prototype. Audio processing runs on your local machine through WebGPU; the Space server only hosts this interface. No audio is uploaded to the Space.</p>
+<h1>SAM Audio Small — Browser inference prototype</h1>
+<p>This is a browser-side prototype. Audio processing runs locally in your browser using either WebGPU (GPU) or WebAssembly (CPU). The Space server only hosts this interface. No audio is uploaded to the Space.</p>
 <div class="panel">
   <h2>1. Input</h2>
   <label for="sam-audio-file">Audio file</label>
@@ -27,6 +27,7 @@ PAGE = r"""
   <label for="sam-prompt">Describe the sound to isolate (English prompt recommended)</label>
   <input id="sam-prompt" type="text" value="a person speaking" placeholder="a person speaking">
   <div class="row">
+    <div style="flex:1;min-width:220px"><label for="sam-provider">Inference device</label><select id="sam-provider" style="box-sizing:border-box;width:100%;padding:10px;border-radius:8px;border:1px solid var(--border-color-primary);background:var(--input-background-fill);color:var(--body-text-color)"><option value="webgpu" selected>GPU — WebGPU</option><option value="wasm">CPU — WebAssembly (WASM)</option></select></div>
     <div style="flex:1;min-width:180px"><label for="sam-steps">ODE steps</label><input id="sam-steps" type="number" min="1" max="16" value="2"></div>
     <div style="flex:2;min-width:220px"><p class="muted">Start with 1–2 steps to validate inference. The reference configuration uses 16 steps and is much slower.</p></div>
   </div>
@@ -47,13 +48,14 @@ PAGE = r"""
     <li>It uses a small number of ODE steps by default for a first end-to-end test. Quality may be poor at 1–2 steps.</li>
     <li>The reference implementation in the repository documents a 48 kHz audio model, T5-base text embeddings, a midpoint ODE solver, and the first/second 128 latent channels as target/residual.</li>
     <li>WebGPU supports only a subset of ONNX operators. A session may load successfully but fail when the first inference runs.</li>
+    <li>CPU mode uses ONNX Runtime WebAssembly (WASM) and may be much slower than WebGPU. Use a short audio clip first.</li>
     <li>First run may download several GB and use substantial browser RAM/VRAM. Use a short audio clip first.</li>
     <li>ONNX files and external weights are cached in this browser when storage quota allows. The browser may evict cached files, and clearing site data removes them.</li>
   </ul>
 </div>
 </div>
 """
-
+ 
 CLIENT_JS = r"""
 async () => {
   const root = document.getElementById("sam-browser-app");
@@ -71,8 +73,8 @@ async () => {
   let ort = null;
   let tokenizer = null;
   let sessions = null;
+  let sessionsProvider = null;
   let previousUrls = [];
-
   function setLog(message) { $("sam-log").textContent = message; }
   function setStatus(message) { $("sam-status").textContent = message; }
   function errorDetails(error) {
@@ -130,14 +132,14 @@ async () => {
     }
     return new Uint8Array(await response.arrayBuffer());
   }
-  async function createSession(filename) {
+  async function createSession(filename, provider) {
     const url = MODEL_BASE + "/" + filename;
     const runtime = await loadOrt();
     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: ["webgpu"],
+      executionProviders: [provider],
       graphOptimizationLevel: "all",
       externalData: [{ path: filename + ".data", data: externalBytes.buffer }]
     });
@@ -148,6 +150,7 @@ async () => {
         try { await session.release(); } catch (_) {}
       }
       sessions = null;
+      sessionsProvider = null;
     }
     if ("caches" in window) {
       try { await caches.delete(CACHE_NAME); } catch (_) {}
@@ -155,8 +158,15 @@ async () => {
     setStatus("Cache cleared");
     setLog("The browser's SAM Audio ONNX cache was cleared. Reload the page before running inference again.");
   }
-  async function loadModels() {
-    if (sessions) return sessions;
+  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"],
@@ -168,11 +178,12 @@ async () => {
       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);
+        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)) {
@@ -301,7 +312,7 @@ async () => {
       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];
+          chunk[c \ chunkSize + t] = latent[c \ timeSteps + start + t];
         }
       }
       const out = await session.run({
@@ -330,7 +341,7 @@ async () => {
     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);
+      view.setInt16(44 + i \ 2, sample < 0 ? sample \ 32768 : sample * 32767, true);
     }
     return new Blob([buffer], {type:"audio/wav"});
   }
@@ -344,13 +355,17 @@ async () => {
     const file = $("sam-audio-file").files?.[0];
     const prompt = $("sam-prompt").value.trim();
     const steps = Number.parseInt($("sam-steps").value, 10);
+    const provider = $("sam-provider").value;
+    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);
     if (audio.samples.length > 48000 * 20) throw new Error("For this first prototype, use audio no longer than 20 seconds.");
-    await loadModels();
-    setStatus("Encoding audio…");
+    setStatus("Loading models for " + deviceLabel + "…");
+    await loadModels(provider);
+    const startedAt = performance.now();
+    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;
@@ -373,16 +388,16 @@ async () => {
       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);
+      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];
+        targetLatent[c \ timeSteps + t] = result[t \ 256 + c];
+        residualLatent[c \ timeSteps + t] = result[t \ 256 + latentChannels + c];
       }
     }
     setStatus("Decoding target…");
@@ -394,6 +409,7 @@ async () => {
     showOutput("sam-residual-player", "sam-residual-download", residual.slice(0, expectedSamples), 48000, "residual.wav");
     setStatus("Finished");
     setLog("Finished browser inference.\nInput: " + file.name + "\nPrompt: " + prompt +
+      "\nDevice: " + deviceLabel + "\nProcessing time: " + ((performance.now() - startedAt) / 1000).toFixed(2) + " s" +
       "\nODE steps: " + steps + "\nInput samples: " + expectedSamples +
       "\nLatent shape: " + JSON.stringify(encoded.dims) +
       "\nTarget samples: " + Math.min(target.length, expectedSamples) +
@@ -417,19 +433,19 @@ async () => {
   });
 }
 """
-
+ 
 @spaces.GPU
 def zerogpu_startup_marker():
     return "No server inference; browser WebGPU prototype."
-
+ 
 with gr.Blocks(title="SAM Audio Browser WebGPU Inference Prototype") as demo:
     gr.Markdown(
-        "## SAM Audio Small — browser WebGPU inference prototype\n"
+        "## SAM Audio Small — browser inference prototype\n"
         "ZeroGPU is selected only because of the Space's hosting configuration. "
         "The marker function is never called. All model inference is attempted locally in the browser."
     )
     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)
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)