vh
a727b93b1e
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.
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
# 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
|
|
from contextlib import asynccontextmanager
|
|
from typing import Optional
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from diffusers import StableAudioPipeline
|
|
from fastapi import FastAPI, HTTPException, Response
|
|
from pydantic import BaseModel, Field
|
|
|
|
MODEL_ID = os.environ.get("SAO_MODEL", "stabilityai/stable-audio-open-1.0")
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
|
|
|
|
state: dict = {}
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
print(f"[sao] loading {MODEL_ID} on {DEVICE} ({DTYPE})", flush=True)
|
|
t0 = time.time()
|
|
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()
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
class SfxRequest(BaseModel):
|
|
prompt: str = Field(..., min_length=1)
|
|
negative_prompt: Optional[str] = "Low quality."
|
|
duration: float = Field(10.0, gt=0.5, le=47.0)
|
|
steps: int = Field(100, ge=10, le=300)
|
|
seed: Optional[int] = None
|
|
cfg_scale: float = Field(7.0, gt=0.0, le=20.0)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model": MODEL_ID,
|
|
"device": DEVICE,
|
|
"loaded": "pipe" in state,
|
|
}
|
|
|
|
|
|
@app.post("/v1/audio/sfx")
|
|
async def sfx(req: SfxRequest):
|
|
pipe = state.get("pipe")
|
|
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)
|
|
|
|
# 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()
|
|
|
|
buf = io.BytesIO()
|
|
sf.write(buf, waveform, pipe.vae.sampling_rate, format="WAV")
|
|
buf.seek(0)
|
|
return Response(content=buf.read(), media_type="audio/wav")
|