diff --git a/app/asr/workers/faster_whisper_worker.py b/app/asr/workers/faster_whisper_worker.py index 4f3f8f6..5f684e1 100644 --- a/app/asr/workers/faster_whisper_worker.py +++ b/app/asr/workers/faster_whisper_worker.py @@ -167,12 +167,23 @@ def _load_model(model_name: str) -> WhisperModel: if _MODEL_CACHE: # Only one whisper model resident at a time — switching sizes frees the old one. _unload_whisper_models() - _MODEL_CACHE[model_name] = WhisperModel( - model_name, - device=DEVICE, - compute_type=COMPUTE_TYPE, - download_root=str(MODEL_CACHE), - ) + try: + _MODEL_CACHE[model_name] = WhisperModel( + model_name, + device=DEVICE, + compute_type=COMPUTE_TYPE, + download_root=str(MODEL_CACHE), + ) + except Exception as e: + # A load interrupted partway (e.g. OOM) can leave CUDA memory + # fragmented/leaked in ways gc.collect()+empty_cache() don't + # reliably reclaim, and since _MODEL_CACHE never got populated + # the idle-unload loop has nothing to clean up either. + # Restarting the whole process is the only guaranteed way to get + # that memory back — supervisord's autorestart=true respawns it + # immediately. + print(f"[faster_whisper] model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}", flush=True) + os._exit(1) return _MODEL_CACHE[model_name] diff --git a/app/asr/workers/qwen3_worker.py b/app/asr/workers/qwen3_worker.py index f36e04b..c15aeef 100644 --- a/app/asr/workers/qwen3_worker.py +++ b/app/asr/workers/qwen3_worker.py @@ -109,15 +109,25 @@ def _load_model(model_id: str) -> Tuple[Any, Any]: _log(f"loading model {model_id}") dtype = torch.float16 if DEVICE == "cuda" else torch.float32 - processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE)) - model = AutoModelForSpeechSeq2Seq.from_pretrained( - model_id, - dtype=dtype, - low_cpu_mem_usage=True, - cache_dir=str(MODEL_CACHE), - ) - model.to(DEVICE) - model.eval() + try: + processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE)) + model = AutoModelForSpeechSeq2Seq.from_pretrained( + model_id, + dtype=dtype, + low_cpu_mem_usage=True, + cache_dir=str(MODEL_CACHE), + ) + model.to(DEVICE) + model.eval() + except Exception as e: + # A load interrupted partway (e.g. OOM) can leave CUDA memory + # fragmented/leaked in ways gc.collect()+empty_cache() don't reliably + # reclaim, and since _MODEL_CACHE never got populated the idle-unload + # loop has nothing to clean up either. Restarting the whole process + # is the only guaranteed way to get that memory back — supervisord's + # autorestart=true respawns it immediately. + _log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}") + os._exit(1) _MODEL_CACHE[model_id] = (model, processor) _log(f"model {model_id} loaded") diff --git a/app/asr/workers/vibevoice_worker.py b/app/asr/workers/vibevoice_worker.py index d46e312..93a373c 100644 --- a/app/asr/workers/vibevoice_worker.py +++ b/app/asr/workers/vibevoice_worker.py @@ -77,20 +77,30 @@ def _load_model(model_id: str) -> Any: from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor _log(f"loading model {model_id}") - processor = VibeVoiceASRProcessor.from_pretrained( - model_id, - language_model_pretrained_name="Qwen/Qwen2.5-1.5B", - cache_dir=str(MODEL_CACHE), - ) - model = VibeVoiceASRForConditionalGeneration.from_pretrained( - model_id, - dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32, - attn_implementation="sdpa", - trust_remote_code=True, - cache_dir=str(MODEL_CACHE), - ) - model.to(DEVICE) - model.eval() + try: + processor = VibeVoiceASRProcessor.from_pretrained( + model_id, + language_model_pretrained_name="Qwen/Qwen2.5-1.5B", + cache_dir=str(MODEL_CACHE), + ) + model = VibeVoiceASRForConditionalGeneration.from_pretrained( + model_id, + dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32, + attn_implementation="sdpa", + trust_remote_code=True, + cache_dir=str(MODEL_CACHE), + ) + model.to(DEVICE) + model.eval() + except Exception as e: + # A load interrupted partway (e.g. OOM) can leave CUDA memory + # fragmented/leaked in ways gc.collect()+empty_cache() don't reliably + # reclaim, and since _MODEL_CACHE never got populated the idle-unload + # loop has nothing to clean up either. Restarting the whole process + # is the only guaranteed way to get that memory back — supervisord's + # autorestart=true respawns it immediately. + _log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}") + os._exit(1) _MODEL_CACHE[model_id] = (model, processor) _log(f"model {model_id} loaded") diff --git a/app/tts/workers/xtts_worker.py b/app/tts/workers/xtts_worker.py index 43cce17..0ff524c 100644 --- a/app/tts/workers/xtts_worker.py +++ b/app/tts/workers/xtts_worker.py @@ -83,7 +83,17 @@ def _load_model() -> Any: from TTS.api import TTS _log(f"loading model {MODEL_NAME}") - tts = TTS(MODEL_NAME).to(_device()) + try: + tts = TTS(MODEL_NAME).to(_device()) + except Exception as e: + # A load interrupted partway (e.g. OOM) can leave CUDA memory + # fragmented/leaked in ways gc.collect()+empty_cache() don't reliably + # reclaim, and since _MODEL_CACHE never got populated the idle-unload + # loop has nothing to clean up either. Restarting the whole process + # is the only guaranteed way to get that memory back — supervisord's + # autorestart=true respawns it immediately. + _log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}") + os._exit(1) _MODEL_CACHE[MODEL_NAME] = tts _log("model loaded") return tts