File size: 1,220 Bytes
e269658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
import whisper
from config import WHISPER_MODEL_SIZE
# Global variables to store models
whisper_processor = None
whisper_model = None
whisper_model_small = None
def load_models():
global whisper_processor, whisper_model, whisper_model_small
if whisper_processor is None:
whisper_processor = WhisperProcessor.from_pretrained(f"openai/whisper-{WHISPER_MODEL_SIZE}")
if whisper_model is None:
whisper_model = WhisperForConditionalGeneration.from_pretrained(f"openai/whisper-{WHISPER_MODEL_SIZE}").to(get_device())
if whisper_model_small is None:
whisper_model_small = whisper.load_model(WHISPER_MODEL_SIZE)
def get_device():
return "cuda:0" if torch.cuda.is_available() else "cpu"
def get_processor():
global whisper_processor
if whisper_processor is None:
load_models()
return whisper_processor
def get_model():
global whisper_model
if whisper_model is None:
load_models()
return whisper_model
def get_whisper_model_small():
global whisper_model_small
if whisper_model_small is None:
load_models()
return whisper_model_small |