admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

cuda or cpu auto detection

Browse source
This commit is contained in:
Quentin Fuxa 2025-02-07 10:16:03 +01:00
parent 4d1aa4421a
commit 0d874fb515
1 changed files with 6 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
src/whisper_streaming/backends.py
View file

@ -4,7 +4,7 @@ import logging
import io import io
import soundfile as sf import soundfile as sf
import math import math
import torch
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -102,11 +102,13 @@ class FasterWhisperASR(ASRBase):
else: else:
raise ValueError("modelsize or model_dir parameter must be set") raise ValueError("modelsize or model_dir parameter must be set")
# this worked fast and reliably on NVIDIA L40 device = "cuda" if torch.cuda.is_available() else "cpu"
compute_type = "float16" if device == "cuda" else "float32"
model = WhisperModel( model = WhisperModel(
model_size_or_path, model_size_or_path,
device="cuda", device=device,
compute_type="float16", compute_type=compute_type,
download_root=cache_dir, download_root=cache_dir,
) )