kjbytes commited on
Commit
187d5d8
·
1 Parent(s): cfff3e1

feat: add thread lock to prevent concurrent model corruption

Browse files

Add threading lock around model.generate() to prevent concurrent calls from corrupting shared model state (alignment_stream_analyzer). Update comments to explain the concurrency issue and upgrade path.

Files changed (1) hide show
  1. app.py +15 -7
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import os
2
  import re
3
  import secrets
 
4
  import uuid
5
 
6
  from fastapi import Depends, FastAPI, HTTPException, Header, Request
@@ -14,6 +15,11 @@ os.makedirs(OUT_DIR, exist_ok=True)
14
 
15
  app = FastAPI(title="Chatterbox TTS")
16
  _model = None
 
 
 
 
 
17
 
18
 
19
  def auth(x_api_key: str = Header(None)):
@@ -23,7 +29,8 @@ def auth(x_api_key: str = Header(None)):
23
 
24
  def model():
25
  # ponytail: lazy singleton so the Space boots fast and healthchecks pass.
26
- # Single uvicorn worker, so no lock needed; add one if we ever scale up.
 
27
  global _model
28
  if _model is None:
29
  import functools
@@ -58,12 +65,13 @@ def tts(req: TTSRequest, request: Request):
58
  if lang not in m.get_supported_languages():
59
  raise HTTPException(400, f"unsupported lang: {lang}")
60
 
61
- wav = m.generate(
62
- req.text,
63
- language_id=lang,
64
- exaggeration=req.exaggeration,
65
- cfg_weight=req.cfg_weight,
66
- )
 
67
  stem = uuid.uuid4().hex
68
  torchaudio.save(os.path.join(OUT_DIR, stem + ".wav"), wav, m.sr)
69
  with open(os.path.join(OUT_DIR, stem + ".srt"), "w") as f:
 
1
  import os
2
  import re
3
  import secrets
4
+ import threading
5
  import uuid
6
 
7
  from fastapi import Depends, FastAPI, HTTPException, Header, Request
 
15
 
16
  app = FastAPI(title="Chatterbox TTS")
17
  _model = None
18
+ # ponytail: global lock. chatterbox keeps per-generation state on the shared model
19
+ # (alignment_stream_analyzer), so concurrent generate() calls corrupt each other.
20
+ # Sync endpoints run in a threadpool, so "one worker" is not one caller.
21
+ # Upgrade path: a real request queue / model pool if throughput matters.
22
+ _gen_lock = threading.Lock()
23
 
24
 
25
  def auth(x_api_key: str = Header(None)):
 
29
 
30
  def model():
31
  # ponytail: lazy singleton so the Space boots fast and healthchecks pass.
32
+ # Worst case two requests both build it and one wins; generate() is serialized
33
+ # by _gen_lock anyway.
34
  global _model
35
  if _model is None:
36
  import functools
 
65
  if lang not in m.get_supported_languages():
66
  raise HTTPException(400, f"unsupported lang: {lang}")
67
 
68
+ with _gen_lock:
69
+ wav = m.generate(
70
+ req.text,
71
+ language_id=lang,
72
+ exaggeration=req.exaggeration,
73
+ cfg_weight=req.cfg_weight,
74
+ )
75
  stem = uuid.uuid4().hex
76
  torchaudio.save(os.path.join(OUT_DIR, stem + ".wav"), wav, m.sr)
77
  with open(os.path.join(OUT_DIR, stem + ".srt"), "w") as f: