local-whisper/local_whisper/inference.py

import os
from threading import Lock

import torch
import whisper

from .settings import whisper_settings

# TODO use pydantic config
model_name = whisper_settings.base_asr_model
if torch.cuda.is_available():
    model = whisper.load_model(model_name).cuda()
else:
    model = whisper.load_model(model_name)
model_lock = Lock()


# TODO move transcribe to a modeling worker
def transcribe(audio):
    # options_dict = {"task" : task}
    # if language:
    #     options_dict["language"] = language
    # if initial_prompt:
    #     options_dict["initial_prompt"] = initial_prompt
    with model_lock:
        # result = model.transcribe(audio, **options_dict)
        result = model.transcribe(audio)

    return result