diff --git a/open_dubbing/dubbing.py b/open_dubbing/dubbing.py index 32cbf99..4cc2a4d 100644 --- a/open_dubbing/dubbing.py +++ b/open_dubbing/dubbing.py @@ -125,6 +125,7 @@ def __init__( self._number_of_steps = number_of_steps self.tts = tts self.device = device + self.cpu_threads = cpu_threads @functools.cached_property def input_file(self): diff --git a/open_dubbing/main.py b/open_dubbing/main.py index fdf2a49..59f09ec 100644 --- a/open_dubbing/main.py +++ b/open_dubbing/main.py @@ -173,7 +173,7 @@ def main(): help=("Device to use"), ) parser.add_argument( - "--cpu-threads", + "--cpu_threads", type=int, default=0, help="number of threads used for CPU inference (if not specified uses defaults for each framework)", diff --git a/open_dubbing/speech_to_text.py b/open_dubbing/speech_to_text.py index 80cefdb..cfc02a7 100644 --- a/open_dubbing/speech_to_text.py +++ b/open_dubbing/speech_to_text.py @@ -20,9 +20,10 @@ class SpeechToText: - def __init__(self, device="cpu"): + def __init__(self, device="cpu", cpu_threads=0): self.model = None self.device = device + self.cpu_threads = cpu_threads logging.getLogger("faster_whisper").setLevel(logging.ERROR) @property @@ -37,6 +38,7 @@ def load_model(self): self._model = WhisperModel( model_size_or_path="large-v3", device=self.device, + cpu_threads=self.cpu_threads, compute_type="float16" if self.device == "cuda" else "int8", ) diff --git a/r.sh b/r.sh index 07629a5..89a0efc 100755 --- a/r.sh +++ b/r.sh @@ -4,7 +4,7 @@ rm -r -f output/ declare -a inputs=("videos/cameratips.mp4" "videos/jordi.mp4" "videos/jobinterview.mp4" "videos/jordimaskudosallhands.mp4" "videos/michael.mp4" "videos/simplicty.mp4") declare -a target_languages=("cat") # Catalan (cat) and French (fra) -declare -a tts_list=("mms") +declare -a tts_list=("coqui") declare -a inputs=("videos/jordimaskudosallhands.mp4") for tts in "${tts_list[@]}"; do @@ -20,7 +20,8 @@ for tts in "${tts_list[@]}"; do --source_language=eng \ --target_language="$language" \ --tts="$tts" \ - --device=cpu + --device=cpu \ + --cpu_threads=8 done done done