unlimited-dubbing-ai / pipeline /asr_diarization.py
ankitranjan52's picture
Update pipeline/asr_diarization.py
da574ab verified
Raw
History Blame Contribute Delete
864 Bytes
import whisperx
_asr_model = None
_alignment_model = None
_metadata = None
def load_asr(device="cpu", model_size="medium"):
global _asr_model
if _asr_model is None:
_asr_model = whisperx.load_model(model_size, device=device, compute_type="float32")
return _asr_model
def transcribe_and_align(audio_path: str, device="cpu"):
model = load_asr(device=device)
result = model.transcribe(audio_path, batch_size=16)
global _alignment_model, _metadata
if _alignment_model is None or _metadata is None:
_alignment_model, _metadata = whisperx.load_align_model(
language_code=result["language"], device=device
)
# Pass device here as required
result_aligned = whisperx.align(
result["segments"], _alignment_model, _metadata, audio_path, device=device
)
return result, result_aligned