From a727b93b1e5f0514f4315bd932c05e517d5df37d Mon Sep 17 00:00:00 2001
From: Vuong Hoang <vh@phasefinal.com>
Date: Sun, 10 May 2026 18:19:24 -0700
Subject: [PATCH] sao: serialize inference under an asyncio.Lock
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

StableAudioPipeline isn't reentrant — concurrent requests share the
scheduler's step_index counter and corrupt each other mid-run
(observed: IndexError in cosine_dpmsolver_multistep when two requests
overlap). Wrap the pipeline call + audio decode in a single
asyncio.Lock created at startup, and run the (sync, GPU-bound)
pipeline call via asyncio.to_thread so the event loop stays
responsive. Concurrent requests now queue cleanly instead of racing.

Verified: 5 parallel POSTs at steps=50 all return 200, clear ~4s
serialization spacing (4, 8, 12, 16, 20s wall time), distinct
output hashes per seed.
---
 stacks/stable-audio-open/server.py | 37 ++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/stacks/stable-audio-open/server.py b/stacks/stable-audio-open/server.py
index 7c24c4e..f62854a 100644
--- a/stacks/stable-audio-open/server.py
+++ b/stacks/stable-audio-open/server.py
@@ -1,6 +1,7 @@
 # FastAPI shim around diffusers' StableAudioPipeline.
 # Single endpoint POST /v1/audio/sfx returns a WAV blob.
 # Model is loaded once on startup and held in process memory.
+import asyncio
 import io
 import os
 import time
@@ -27,6 +28,11 @@ async def lifespan(app: FastAPI):
     pipe = StableAudioPipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE)
     pipe = pipe.to(DEVICE)
     state["pipe"] = pipe
+    # Single-GPU diffusers pipelines aren't reentrant — concurrent calls
+    # share the same scheduler step_index counter and corrupt each
+    # other (observed: IndexError in cosine_dpmsolver_multistep when
+    # two requests overlap mid-run). Serialize at the request boundary.
+    state["lock"] = asyncio.Lock()
     print(f"[sao] loaded in {time.time() - t0:.1f}s", flush=True)
     yield
     state.clear()
@@ -55,26 +61,33 @@ def health():
 
 
 @app.post("/v1/audio/sfx")
-def sfx(req: SfxRequest):
+async def sfx(req: SfxRequest):
     pipe = state.get("pipe")
-    if pipe is None:
+    lock = state.get("lock")
+    if pipe is None or lock is None:
         raise HTTPException(503, "model not loaded yet")
 
     generator = None
     if req.seed is not None:
         generator = torch.Generator(DEVICE).manual_seed(req.seed)
 
-    audio = pipe(
-        req.prompt,
-        negative_prompt=req.negative_prompt,
-        num_inference_steps=req.steps,
-        audio_end_in_s=req.duration,
-        num_waveforms_per_prompt=1,
-        guidance_scale=req.cfg_scale,
-        generator=generator,
-    ).audios
+    # Hold the lock for the whole inference + encode. Concurrent
+    # requests queue cleanly instead of racing the scheduler's
+    # step_index counter into an IndexError.
+    async with lock:
+        audio = await asyncio.to_thread(
+            lambda: pipe(
+                req.prompt,
+                negative_prompt=req.negative_prompt,
+                num_inference_steps=req.steps,
+                audio_end_in_s=req.duration,
+                num_waveforms_per_prompt=1,
+                guidance_scale=req.cfg_scale,
+                generator=generator,
+            ).audios
+        )
+        waveform = audio[0].T.float().cpu().numpy()
 
-    waveform = audio[0].T.float().cpu().numpy()
     buf = io.BytesIO()
     sf.write(buf, waveform, pipe.vae.sampling_rate, format="WAV")
     buf.seek(0)