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

src-python/config.py

@@ -724,6 +724,28 @@ class Config:
             self._USE_WHISPER_FEATURE = value
             saveJson(self.PATH_CONFIG, inspect.currentframe().f_code.co_name, value)
 
+    @property
+    @json_serializable('SELECTED_TRANSLATION_COMPUTE_DEVICE')
+    def SELECTED_TRANSLATION_COMPUTE_DEVICE(self):
+        return self._SELECTED_TRANSLATION_COMPUTE_DEVICE
+
+    @SELECTED_TRANSLATION_COMPUTE_DEVICE.setter
+    def SELECTED_TRANSLATION_COMPUTE_DEVICE(self, value):
+        if isinstance(value, dict):
+            self._SELECTED_TRANSLATION_COMPUTE_DEVICE = value
+            saveJson(self.PATH_CONFIG, inspect.currentframe().f_code.co_name, value)
+
+    @property
+    @json_serializable('SELECTED_TRANSCRIPTION_COMPUTE_DEVICE')
+    def SELECTED_TRANSCRIPTION_COMPUTE_DEVICE(self):
+        return self._SELECTED_TRANSCRIPTION_COMPUTE_DEVICE
+
+    @SELECTED_TRANSCRIPTION_COMPUTE_DEVICE.setter
+    def SELECTED_TRANSCRIPTION_COMPUTE_DEVICE(self, value):
+        if isinstance(value, dict):
+            self._SELECTED_TRANSCRIPTION_COMPUTE_DEVICE = value
+            saveJson(self.PATH_CONFIG, inspect.currentframe().f_code.co_name, value)
+
     @property
     @json_serializable('CTRANSLATE2_WEIGHT_TYPE')
     def CTRANSLATE2_WEIGHT_TYPE(self):
@@ -1105,8 +1127,10 @@ class Config:
         }
         self._USE_EXCLUDE_WORDS = True
         self._USE_TRANSLATION_FEATURE = True
-        self._CTRANSLATE2_WEIGHT_TYPE = "Small"
         self._USE_WHISPER_FEATURE = False
+        self._SELECTED_TRANSLATION_COMPUTE_DEVICE = {"type": "cpu", "index": 0, "name":"cpu"}
+        self._SELECTED_TRANSCRIPTION_COMPUTE_DEVICE = {"type": "cpu", "index": 0, "name":"cpu"}
+        self._CTRANSLATE2_WEIGHT_TYPE = "Small"
         self._WHISPER_WEIGHT_TYPE = "base"
         self._SEND_MESSAGE_FORMAT = "[message]"
         self._SEND_MESSAGE_FORMAT_WITH_T = "[message]([translation])"

src-python/model.py

@@ -109,7 +109,11 @@ class Model:
         return checkCTranslate2Weight(config.PATH_LOCAL, config.CTRANSLATE2_WEIGHT_TYPE)
 
     def changeTranslatorCTranslate2Model(self):
-        self.translator.changeCTranslate2Model(config.PATH_LOCAL, config.CTRANSLATE2_WEIGHT_TYPE)
+        self.translator.changeCTranslate2Model(
+            config.PATH_LOCAL,
+            config.CTRANSLATE2_WEIGHT_TYPE,
+            config.SELECTED_TRANSLATION_COMPUTE_DEVICE["type"],
+            config.SELECTED_TRANSLATION_COMPUTE_DEVICE["device_index"])
 
     def downloadCTranslate2ModelWeight(self, callbackFunc=None):
         return downloadCTranslate2Weight(config.PATH_LOCAL, config.CTRANSLATE2_WEIGHT_TYPE, callbackFunc)
@@ -425,6 +429,8 @@ class Model:
             transcription_engine=config.SELECTED_TRANSCRIPTION_ENGINE,
             root=config.PATH_LOCAL,
             whisper_weight_type=config.WHISPER_WEIGHT_TYPE,
+            device=config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE["type"],
+            device_index=config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE["device_index"],
         )
         def sendMicTranscript():
             try:
@@ -587,6 +593,8 @@ class Model:
             transcription_engine=config.SELECTED_TRANSCRIPTION_ENGINE,
             root=config.PATH_LOCAL,
             whisper_weight_type=config.WHISPER_WEIGHT_TYPE,
+            device=config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE["type"],
+            device_index=config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE["device_index"],
         )
         def sendSpeakerTranscript():
             try:

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=device,
-        device_index=0,
+        device_index=device_index,
-        compute_type="int8",
+        compute_type=compute_type,
         cpu_threads=4,
         num_workers=1,
         local_files_only=True,

src-python/models/translation/translation_translator.py

@@ -29,17 +29,19 @@ class Translator():
             result = False
         return result
 
-    def changeCTranslate2Model(self, path, model_type):
+    def changeCTranslate2Model(self, path, model_type, device="cpu", device_index=0):
         self.is_loaded_ctranslate2_model = False
         directory_name = ctranslate2_weights[model_type]["directory_name"]
         tokenizer = ctranslate2_weights[model_type]["tokenizer"]
         weight_path = os_path.join(path, "weights", "ctranslate2", directory_name)
         tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer")
+
+        compute_type = "int8" if device == "cpu" else "float16"
         self.ctranslate2_translator = ctranslate2.Translator(
             weight_path,
-            device="cuda",
+            device=device,
-            device_index=0,
+            device_index=device_index,
-            compute_type="int8",
+            compute_type=compute_type,
             inter_threads=1,
             intra_threads=4
         )

src-python/webui_controller.py

@@ -7,6 +7,7 @@ from device_manager import device_manager
 from config import config
 from model import model
 from utils import isUniqueStrings, printLog
+import torch
 
 class Controller:
     def __init__(self) -> None:
@@ -363,10 +364,36 @@ class Controller:
     def getMessageBoxRatioRange(*args, **kwargs) -> dict:
         return {"status":200, "result":config.MESSAGE_BOX_RATIO_RANGE}
 
+    @staticmethod
+    def getComputeDeviceList(*args, **kwargs) -> dict:
+        device_list = [{"type":"cuda", "device_index": i, "name": torch.cuda.get_device_name(i)} for i in range(torch.cuda.device_count())]
+        device_list.append({"type":"cpu", "device_index": 0, "name": "cpu"})
+        return {"status":200, "result":device_list}
+
+    @staticmethod
+    def getSelectedTranslationComputeDevice(*args, **kwargs) -> dict:
+        return {"status":200, "result":config.SELECTED_TRANSLATION_COMPUTE_DEVICE}
+
+    @staticmethod
+    def setSelectedTranslationComputeDevice(device:str, *args, **kwargs) -> dict:
+        printLog("setSelectedTranslationComputeDevice", device)
+        config.SELECTED_TRANSLATION_COMPUTE_DEVICE = device
+        return {"status":200,"result":config.SELECTED_TRANSLATION_COMPUTE_DEVICE}
+
     @staticmethod
     def getSelectableCtranslate2WeightTypeDict(*args, **kwargs) -> dict:
         return {"status":200, "result":config.SELECTABLE_CTRANSLATE2_WEIGHT_TYPE_DICT}
 
+    @staticmethod
+    def getSelectedTranscriptionComputeDevice(*args, **kwargs) -> dict:
+        return {"status":200, "result":config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE}
+
+    @staticmethod
+    def setSelectedTranscriptionComputeDevice(device:str, *args, **kwargs) -> dict:
+        printLog("setSelectedTranscriptionComputeDevice", device)
+        config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE = device
+        return {"status":200,"result":config.SELECTED_TRANSCRIPTION_COMPUTE_DEVICE}
+
     @staticmethod
     def getSelectableWhisperModelTypeDict(*args, **kwargs) -> dict:
         return {"status":200, "result":config.SELECTABLE_WHISPER_WEIGHT_TYPE_DICT}

src-python/webui_mainloop.py

@@ -128,6 +128,10 @@ mapping = {
     "/set/enable/use_translation_feature": {"status": True, "variable":controller.setEnableUseTranslationFeature},
     "/set/disable/use_translation_feature": {"status": True, "variable":controller.setDisableUseTranslationFeature},
 
+    "/get/data/translation_compute_device_dict": {"status": True, "variable":controller.getComputeDeviceList},
+    "/get/data/selected_translation_compute_device": {"status": True, "variable":controller.getSelectedTranslationComputeDevice},
+    "/set/data/selected_translation_compute_device": {"status": True, "variable":controller.setSelectedTranslationComputeDevice},
+
     "/get/data/selectable_ctranslate2_weight_type_dict": {"status": True, "variable":controller.getSelectableCtranslate2WeightTypeDict},
 
     "/get/data/ctranslate2_weight_type": {"status": True, "variable":controller.getCtranslate2WeightType},
@@ -229,6 +233,10 @@ mapping = {
     "/set/enable/check_speaker_threshold": {"status": True, "variable":controller.setEnableCheckSpeakerThreshold},
     "/set/disable/check_speaker_threshold": {"status": True, "variable":controller.setDisableCheckSpeakerThreshold},
 
+    "/get/data/transcription_compute_device_dict": {"status": True, "variable":controller.getComputeDeviceList},
+    "/get/data/selected_transcription_compute_device": {"status": True, "variable":controller.getSelectedTranscriptionComputeDevice},
+    "/set/data/selected_transcription_compute_device": {"status": True, "variable":controller.setSelectedTranscriptionComputeDevice},
+
     "/get/data/selectable_whisper_weight_type_dict": {"status": True, "variable":controller.getSelectableWhisperModelTypeDict},
 
     "/get/data/whisper_weight_type": {"status": True, "variable":controller.getWhisperWeightType},