diff --git a/buzz/file_transcriber_queue_worker.py b/buzz/file_transcriber_queue_worker.py index e4b61b28..72db8893 100644 --- a/buzz/file_transcriber_queue_worker.py +++ b/buzz/file_transcriber_queue_worker.py @@ -134,75 +134,79 @@ class FileTranscriberQueueWorker(QObject): logging.debug("Waiting for next transcription task") - # Clean up of previous run. + self._cleanup_previous_transcriber() + + if not self._get_next_task(): + self.is_running = False + self.completed.emit() + return + + self.is_running = True + + if self.current_task.transcription_options.extract_speech: + status = self._setup_speech_extraction() + if status == "error": + self.is_running = False + return + + self._run_plugins() + + logging.debug("Starting next transcription task") + self.task_progress.emit(self.current_task, 0) + + self._create_transcriber() + self._setup_transcriber_thread() + + def _cleanup_previous_transcriber(self): if self.current_transcriber is not None: self.current_transcriber.stop() self.current_transcriber = None - # Get next non-canceled task from queue + def _get_next_task(self) -> bool: while True: - self.current_task: Optional[FileTranscriptionTask] = self.tasks_queue.get() - - # Stop listening when a "None" task is received + self.current_task = self.tasks_queue.get() if self.current_task is None: - self.is_running = False - self.completed.emit() - return - + return False if self.current_task.uid in self.canceled_tasks: continue + return True - break + def _setup_speech_extraction(self) -> str: + logging.debug("Will extract speech") - # Set is_running AFTER we have a valid task to process - self.is_running = True + force_cpu = os.getenv("BUZZ_FORCE_CPU", "false").lower() == "true" + if force_cpu: + device = "cpu" + else: + import torch + device = "cuda" if torch.cuda.is_available() else "cpu" - if self.current_task.transcription_options.extract_speech: - logging.debug("Will extract speech") + task_file_path = Path(self.current_task.file_path) + speech_path = task_file_path.with_name(f"{task_file_path.stem}_speech.mp3") - # Force CPU if specified, otherwise use CUDA if available - force_cpu = os.getenv("BUZZ_FORCE_CPU", "false").lower() == "true" - if force_cpu: - device = "cpu" - else: - import torch - device = "cuda" if torch.cuda.is_available() else "cpu" + status = self._extract_speech(str(task_file_path), str(speech_path), device) - task_file_path = Path(self.current_task.file_path) - speech_path = task_file_path.with_name(f"{task_file_path.stem}_speech.mp3") + if status == "error": + self.task_error.emit( + self.current_task, + _("Speech extraction failed! Check your internet connection \u2014 a model may need to be downloaded."), + ) + elif status == "ok": + self.speech_path = speech_path + if not self.current_task.original_file_path: + self.current_task.original_file_path = str(task_file_path) + self.current_task.file_path = str(speech_path) - status = self._extract_speech(str(task_file_path), str(speech_path), device) + return status - if status == "error": - self.task_error.emit( - self.current_task, - _("Speech extraction failed! Check your internet connection — a model may need to be downloaded."), - ) - self.is_running = False - return - - if status == "ok": - self.speech_path = speech_path - # Remember the original audio path: file_path is about to point - # at the temporary "_speech.mp3", which is deleted once the - # transcription completes. Plugins (e.g. the transcript resizer) - # need the original file in their post-completion hooks. - if not self.current_task.original_file_path: - self.current_task.original_file_path = str(task_file_path) - self.current_task.file_path = str(speech_path) - # status == "no_audio": transcribe the original file as-is. - - # Let plugins process / replace the source audio before transcription. - # Runs on this worker thread; plugins may overwrite current_task.file_path. + def _run_plugins(self): if self.plugin_manager is not None: try: self.plugin_manager.run_before_transcription(self.current_task) except Exception as e: logging.error(f"Plugin before_transcription failed: {e}", exc_info=True) - logging.debug("Starting next transcription task") - self.task_progress.emit(self.current_task, 0) - + def _create_transcriber(self): model_type = self.current_task.transcription_options.model.model_type if model_type == ModelType.OPEN_AI_WHISPER_API: self.current_transcriber = OpenAIWhisperAPIFileTranscriber( @@ -218,6 +222,7 @@ class FileTranscriberQueueWorker(QObject): else: raise Exception(f"Unknown model type: {model_type}") + def _setup_transcriber_thread(self): self.current_transcriber_thread = QThread(self) self.current_transcriber.moveToThread(self.current_transcriber_thread) @@ -240,7 +245,6 @@ class FileTranscriberQueueWorker(QObject): self.current_transcriber.completed.connect(self.on_task_completed) - # Wait for next item on the queue self.current_transcriber.error.connect(lambda: self._on_task_finished()) self.current_transcriber.completed.connect(lambda: self._on_task_finished()) diff --git a/buzz/model_loader.py b/buzz/model_loader.py index 3921a7ca..7f0968fb 100644 --- a/buzz/model_loader.py +++ b/buzz/model_loader.py @@ -649,133 +649,135 @@ class ModelDownloader(QRunnable): def _register_process(self, proc: multiprocessing.Process): self._download_process = proc + def _download_whisper_cpp(self) -> None: + if self.custom_model_url: + url = self.custom_model_url + file_path = get_whisper_cpp_file_path( + size=self.model.whisper_model_size) + self.download_model_to_path(url=url, file_path=file_path) + return + + repo_id = WHISPER_CPP_REPO_ID + + if self.model.whisper_model_size == WhisperModelSize.LUMII: + repo_id = WHISPER_CPP_LUMII_REPO_ID + + model_name = self.model.whisper_model_size.to_whisper_cpp_model_size() + + whisper_cpp_model_files = [ + f"ggml-{model_name}.bin", + "README.md" + ] + if self.is_coreml_supported: + whisper_cpp_model_files = [ + f"ggml-{model_name}.bin", + f"ggml-{model_name}-encoder.mlmodelc.zip", + "README.md" + ] + + model_path = download_from_huggingface( + repo_id=repo_id, + allow_patterns=whisper_cpp_model_files, + progress=self.signals.progress, + on_process=self._register_process, + ) + + if self.stopped: + return + + if self.is_coreml_supported: + import tempfile + + target_dir = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc") + zip_path = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc.zip") + + if os.path.exists(target_dir): + shutil.rmtree(target_dir) + + with tempfile.TemporaryDirectory() as temp_dir: + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + zip_ref.extractall(temp_dir) + + macosx_path = os.path.join(temp_dir, "__MACOSX") + if os.path.exists(macosx_path): + shutil.rmtree(macosx_path) + + temp_contents = os.listdir(temp_dir) + if len(temp_contents) == 1 and os.path.isdir(os.path.join(temp_dir, temp_contents[0])): + nested_dir = os.path.join(temp_dir, temp_contents[0]) + shutil.move(nested_dir, target_dir) + else: + os.makedirs(target_dir, exist_ok=True) + for item in temp_contents: + src = os.path.join(temp_dir, item) + dst = os.path.join(target_dir, item) + if os.path.isdir(src): + shutil.copytree(src, dst) + else: + shutil.copy2(src, dst) + + self.signals.finished.emit(os.path.join( + model_path, f"ggml-{model_name}.bin")) + + def _download_whisper(self) -> None: + url = whisper._MODELS[self.model.whisper_model_size.value] + file_path = get_whisper_file_path( + size=self.model.whisper_model_size) + expected_sha256 = url.split("/")[-2] + self.download_model_to_path( + url=url, file_path=file_path, expected_sha256=expected_sha256 + ) + + def _download_faster_whisper(self) -> None: + model_path = download_faster_whisper_model( + model=self.model, + progress=self.signals.progress, + on_process=self._register_process, + ) + + if self.stopped: + return + + if model_path == "": + self.signals.error.emit(_("Error")) + + self.signals.finished.emit(model_path) + + def _download_hugging_face(self) -> None: + model_path = download_from_huggingface( + self.model.hugging_face_model_id, + allow_patterns=HUGGING_FACE_MODEL_ALLOW_PATTERNS, + progress=self.signals.progress, + on_process=self._register_process, + ) + + if self.stopped: + return + + if model_path == "": + self.signals.error.emit(_("Error")) + + self.signals.finished.emit(model_path) + + def _download_openai_whisper_api(self) -> None: + self.signals.finished.emit("") + def run(self) -> None: logging.debug("Downloading model: %s, %s", self.model, self.model.hugging_face_model_id) if self.model.model_type == ModelType.WHISPER_CPP: - if self.custom_model_url: - url = self.custom_model_url - file_path = get_whisper_cpp_file_path( - size=self.model.whisper_model_size) - return self.download_model_to_path(url=url, file_path=file_path) - - repo_id = WHISPER_CPP_REPO_ID - - if self.model.whisper_model_size == WhisperModelSize.LUMII: - repo_id = WHISPER_CPP_LUMII_REPO_ID - - model_name = self.model.whisper_model_size.to_whisper_cpp_model_size() - - whisper_cpp_model_files = [ - f"ggml-{model_name}.bin", - "README.md" - ] - if self.is_coreml_supported: - whisper_cpp_model_files = [ - f"ggml-{model_name}.bin", - f"ggml-{model_name}-encoder.mlmodelc.zip", - "README.md" - ] - - model_path = download_from_huggingface( - repo_id=repo_id, - allow_patterns=whisper_cpp_model_files, - progress=self.signals.progress, - on_process=self._register_process, - ) - - if self.stopped: - return - - if self.is_coreml_supported: - import tempfile - - target_dir = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc") - zip_path = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc.zip") - - # Remove target directory if it exists - if os.path.exists(target_dir): - shutil.rmtree(target_dir) - - # Extract to a temporary directory first - with tempfile.TemporaryDirectory() as temp_dir: - with zipfile.ZipFile(zip_path, 'r') as zip_ref: - zip_ref.extractall(temp_dir) - - # Remove __MACOSX metadata folders if present - macosx_path = os.path.join(temp_dir, "__MACOSX") - if os.path.exists(macosx_path): - shutil.rmtree(macosx_path) - - # Check if there's a single top-level directory - temp_contents = os.listdir(temp_dir) - if len(temp_contents) == 1 and os.path.isdir(os.path.join(temp_dir, temp_contents[0])): - # Single directory - move its contents to target - nested_dir = os.path.join(temp_dir, temp_contents[0]) - shutil.move(nested_dir, target_dir) - else: - # Multiple items or files - copy everything to target - os.makedirs(target_dir, exist_ok=True) - for item in temp_contents: - src = os.path.join(temp_dir, item) - dst = os.path.join(target_dir, item) - if os.path.isdir(src): - shutil.copytree(src, dst) - else: - shutil.copy2(src, dst) - - self.signals.finished.emit(os.path.join( - model_path, f"ggml-{model_name}.bin")) - return - - if self.model.model_type == ModelType.WHISPER: - url = whisper._MODELS[self.model.whisper_model_size.value] - file_path = get_whisper_file_path( - size=self.model.whisper_model_size) - expected_sha256 = url.split("/")[-2] - return self.download_model_to_path( - url=url, file_path=file_path, expected_sha256=expected_sha256 - ) - - if self.model.model_type == ModelType.FASTER_WHISPER: - model_path = download_faster_whisper_model( - model=self.model, - progress=self.signals.progress, - on_process=self._register_process, - ) - - if self.stopped: - return - - if model_path == "": - self.signals.error.emit(_("Error")) - - self.signals.finished.emit(model_path) - return - - if self.model.model_type == ModelType.HUGGING_FACE: - model_path = download_from_huggingface( - self.model.hugging_face_model_id, - allow_patterns=HUGGING_FACE_MODEL_ALLOW_PATTERNS, - progress=self.signals.progress, - on_process=self._register_process, - ) - - if self.stopped: - return - - if model_path == "": - self.signals.error.emit(_("Error")) - - self.signals.finished.emit(model_path) - return - - if self.model.model_type == ModelType.OPEN_AI_WHISPER_API: - self.signals.finished.emit("") - return - - raise Exception("Invalid model type: " + self.model.model_type.value) + self._download_whisper_cpp() + elif self.model.model_type == ModelType.WHISPER: + self._download_whisper() + elif self.model.model_type == ModelType.FASTER_WHISPER: + self._download_faster_whisper() + elif self.model.model_type == ModelType.HUGGING_FACE: + self._download_hugging_face() + elif self.model.model_type == ModelType.OPEN_AI_WHISPER_API: + self._download_openai_whisper_api() + else: + raise Exception("Invalid model type: " + self.model.model_type.value) def download_model_to_path( self, url: str, file_path: str, expected_sha256: Optional[str] = None