audio: FP16 wav2vec2 on GPU (halves VRAM)

Load the wav2vec2/HF-ASR pipeline in FP16 on CUDA (~3GB -> ~1.5GB), FP32 on CPU,
with a fallback to FP32 if a checkpoint rejects half precision; overridable via a
`dtype` config field. Lower the wav2vec2 VRAM reservation default 3->2GB to match.
Enables co-residency of the full transcription/diarization/speaker stack with a
14B LLM on one 24GB 3090 (measured ~14.5GB used with a gemma-4-14b @8k sibling).
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01R6iugoNeEgt9PkajKrFyq4
parent ed8c014d
...@@ -124,17 +124,30 @@ class _Wav2Vec2Backend: ...@@ -124,17 +124,30 @@ class _Wav2Vec2Backend:
return return
import torch import torch
from transformers import pipeline from transformers import pipeline
device = 0 if torch.cuda.is_available() else -1 use_cuda = torch.cuda.is_available()
device = 0 if use_cuda else -1
target = self.config.get("model_path") or self.model_name target = self.config.get("model_path") or self.model_name
# chunk_length_s lets the pipeline transcribe arbitrarily long audio and # Load in FP16 on GPU to roughly halve VRAM (~3GB→~1.5GB) — CTC inference is
# emit per-chunk timestamps for segments. # fine in half precision. CPU stays FP32. Override with config dtype
self._pipe = pipeline( # (float16 | bfloat16 | float32).
"automatic-speech-recognition", pipe_kwargs = dict(
model=target, model=target,
device=device, device=device,
chunk_length_s=int(self.config.get("chunk_length_s", 30) or 30), chunk_length_s=int(self.config.get("chunk_length_s", 30) or 30),
stride_length_s=int(self.config.get("stride_length_s", 5) or 5), stride_length_s=int(self.config.get("stride_length_s", 5) or 5),
) )
if use_cuda:
dt = str(self.config.get("dtype") or "float16").lower()
dtype = {"float16": torch.float16, "fp16": torch.float16, "half": torch.float16,
"bfloat16": torch.bfloat16, "bf16": torch.bfloat16}.get(dt)
if dtype is not None:
pipe_kwargs["torch_dtype"] = dtype
try:
self._pipe = pipeline("automatic-speech-recognition", **pipe_kwargs)
except Exception:
# Some checkpoints/ops reject half precision — fall back to FP32.
pipe_kwargs.pop("torch_dtype", None)
self._pipe = pipeline("automatic-speech-recognition", **pipe_kwargs)
def transcribe(self, audio_path: str, language: Optional[str] = None, def transcribe(self, audio_path: str, language: Optional[str] = None,
prompt: Optional[str] = None, temperature: float = 0.0, prompt: Optional[str] = None, temperature: float = 0.0,
......
...@@ -600,7 +600,7 @@ async def _run_transcription_inner( ...@@ -600,7 +600,7 @@ async def _run_transcription_inner(
uses_gpu = bool(torch.cuda.is_available()) uses_gpu = bool(torch.cuda.is_available())
except Exception: except Exception:
uses_gpu = False uses_gpu = False
needed_gb = float(stt_cfg.get('used_vram_gb') or 3.0) if uses_gpu else 0.0 needed_gb = float(stt_cfg.get('used_vram_gb') or 2.0) if uses_gpu else 0.0
def _run_backend(): def _run_backend():
def _loader(): def _loader():
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment