Add cpu_threads support (#2)

OpenVoiceOS · May 30, 2023 · 15e4d3d · 15e4d3d
1 parent 4251e8e
commit 15e4d3d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -24,10 +24,10 @@ To use Whisper as STT
         "model": "large-v2",
         "use_cuda": true,
         "compute_type": "float16",
-        "beam_size": 5
+        "beam_size": 5,
+        "cpu_threads": 4
     }
   }
-
 ```
 
 To use Whisper for lang detection  (ovos-dinkum-listener only)
@@ -41,7 +41,6 @@ To use Whisper for lang detection  (ovos-dinkum-listener only)
         }
     }
   }
-
 ```
 
 ## Models

diff --git a/ovos_stt_plugin_fasterwhisper/__init__.py b/ovos_stt_plugin_fasterwhisper/__init__.py
@@ -21,6 +21,7 @@ def __init__(self, config=None):
         self.compute_type = self.config.get("compute_type", "int8")
         self.use_cuda = self.config.get("use_cuda", False)
         self.beam_size = self.config.get("beam_size", 5)
+        self.cpu_threads = self.config.get("cpu_threads", 4)
 
         if self.use_cuda:
             device = "cuda"
@@ -187,12 +188,13 @@ def __init__(self, *args, **kwargs):
         self.beam_size = self.config.get("beam_size", 5)
         self.compute_type = self.config.get("compute_type", "int8")
         self.use_cuda = self.config.get("use_cuda", False)
+        self.cpu_threads = self.config.get("cpu_threads", 4)
 
         if self.use_cuda:
             device = "cuda"
         else:
             device = "cpu"
-        self.engine = WhisperModel(model, device=device, compute_type=self.compute_type)
+        self.engine = WhisperModel(model, device=device, compute_type=self.compute_type, cpu_threads=self.cpu_threads)
 
     @staticmethod
     def audiodata2array(audio_data):