sao: serialize inference under an asyncio.Lock

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.
This commit is contained in:
vh
2026-05-10 18:19:24 -07:00
parent 17b9adf29c
commit a727b93b1e
+25 -12
View File
@@ -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)