You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
30 lines
722 B
Python
30 lines
722 B
Python
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
|