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

src-python/models/transcription/transcription_transcriber.py

@@ -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):

src-python/models/transcription/transcription_whisper.py

@@ -1,6 +1,7 @@
 from os import path as os_path, makedirs as os_makedirs
 from requests import get as requests_get
 from typing import Callable
+import torch
 import huggingface_hub
 from faster_whisper import WhisperModel
 import logging
@@ -51,7 +52,7 @@ def checkWhisperWeight(root, weight_type):
     try:
         WhisperModel(
             path,
-            device="cuda",
+            device="cpu",
             device_index=0,
             compute_type="int8",
             cpu_threads=4,
@@ -75,13 +76,14 @@ def downloadWhisperWeight(root, weight_type, callbackFunc):
         url = huggingface_hub.hf_hub_url(_MODELS[weight_type], filename)
         downloadFile(url, file_path, func=callbackFunc)
 
-def getWhisperModel(root, weight_type):
+def getWhisperModel(root, weight_type, device="cpu", device_index=0):
     path = os_path.join(root, "weights", "whisper", weight_type)
+    compute_type = "int8" if device == "cpu" else "float16"
     return WhisperModel(
         path,
-        device="cuda",
-        device_index=0,
-        compute_type="int8",
+        device=device,
+        device_index=device_index,
+        compute_type=compute_type,
         cpu_threads=4,
         num_workers=1,
         local_files_only=True,