👍️[Update] Model: cpu/cudaをtranslationもしくはtranscriptionで選択できるように実装

This commit is contained in:
misyaguziya
2024-10-23 13:41:34 +09:00
parent 2136865493
commit af3fe1f0f9
7 changed files with 84 additions and 13 deletions

View File

@@ -18,7 +18,7 @@ PHRASE_TIMEOUT = 3
MAX_PHRASES = 10
class AudioTranscriber:
def __init__(self, speaker, source, phrase_timeout, max_phrases, transcription_engine, root=None, whisper_weight_type=None):
def __init__(self, speaker, source, phrase_timeout, max_phrases, transcription_engine, root=None, whisper_weight_type=None, device="cpu", device_index=0):
self.speaker = speaker
self.phrase_timeout = phrase_timeout
self.max_phrases = max_phrases
@@ -38,7 +38,7 @@ class AudioTranscriber:
}
if transcription_engine == "Whisper" and checkWhisperWeight(root, whisper_weight_type) is True:
self.whisper_model = getWhisperModel(root, whisper_weight_type)
self.whisper_model = getWhisperModel(root, whisper_weight_type, device=device, device_index=device_index)
self.transcription_engine = "Whisper"
def transcribeAudioQueue(self, audio_queue, language, country, avg_logprob=-0.8, no_speech_prob=0.6):