buzz/buzz/transcriber/whisper_file_transcriber.py

import datetime
import json
import logging
import multiprocessing
import re
import os
import sys

# Preload CUDA libraries before importing torch - required for subprocess contexts
from buzz import cuda_setup  # noqa: F401

import torch
import platform

from buzz.transcriber.cuda_device import cuda_works
import subprocess
from platformdirs import user_cache_dir
from multiprocessing.connection import Connection
from threading import Lock, Thread
from typing import Optional, List

import tqdm
import psutil
from PyQt6.QtCore import QObject

from buzz import whisper_audio
from buzz.conn import pipe_stderr
from buzz.model_loader import ModelType, map_language_to_mms
from buzz.transformers_whisper import TransformersTranscriber
from buzz.transcriber.file_transcriber import FileTranscriber
from buzz.transcriber.transcriber import FileTranscriptionTask, Segment, Task, DEFAULT_WHISPER_TEMPERATURE
from buzz.transcriber.whisper_cpp import WhisperCpp

import av
import faster_whisper
import whisper
import stable_whisper
from stable_whisper import WhisperResult

PROGRESS_REGEX = re.compile(r"\d+(\.\d+)?%")


def terminate_child_processes(pid: int, timeout: float = 5.0) -> None:
    """Terminate every descendant process of ``pid`` (but not ``pid`` itself).

    For whisper.cpp the actual work runs in a ``whisper-cli`` subprocess spawned
    by our multiprocessing worker. ``multiprocessing.Process.terminate()`` only
    signals the worker, leaving that CLI orphaned and still consuming CPU/GPU
    after the user presses Stop or closes the app. This kills those descendants.

    Crucially it does *not* touch ``pid`` itself: the worker is a direct child
    of the app process and must be reaped via ``multiprocessing`` (join). Calling
    ``psutil.wait_procs`` on it here would steal the ``waitpid`` reap and leave
    ``Process.is_alive()`` stuck reporting True. The descendants are grandchildren
    of the app, so psutil only polls them and no such race occurs.
    """
    try:
        parent = psutil.Process(pid)
    except (psutil.NoSuchProcess, ValueError):
        return

    # Snapshot the tree before anything gets reparented by the worker exiting.
    try:
        descendants = parent.children(recursive=True)
    except (psutil.NoSuchProcess, psutil.AccessDenied):
        return

    # Ask each descendant to exit (SIGTERM / TerminateProcess).
    for proc in descendants:
        try:
            proc.terminate()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass

    _, alive = psutil.wait_procs(descendants, timeout=timeout)

    # Force-kill whatever ignored the polite request.
    for proc in alive:
        try:
            proc.kill()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    psutil.wait_procs(alive, timeout=timeout)


def check_file_has_audio_stream(file_path: str) -> None:
    """Check if a media file has at least one audio stream.

    Raises:
        ValueError: If the file has no audio streams.
    """
    try:
        with av.open(file_path) as container:
            if len(container.streams.audio) == 0:
                raise ValueError("No audio streams found")
    except av.error.InvalidDataError as e:
        raise ValueError(f"Invalid media file: {e}")
    except (av.error.FileNotFoundError, OSError, UnicodeDecodeError):
        raise ValueError("File not found")


class WhisperFileTranscriber(FileTranscriber):
    """WhisperFileTranscriber transcribes an audio file to text, writes the text to a file, and then opens the file
    using the default program for opening txt files."""

    current_process: multiprocessing.Process
    running = False
    read_line_thread: Optional[Thread] = None
    READ_LINE_THREAD_STOP_TOKEN = "--STOP--"

    def __init__(
        self, task: FileTranscriptionTask, parent: Optional["QObject"] = None
    ) -> None:
        super().__init__(task, parent)
        self.segments = []
        self.started_process = False
        self.stopped = False
        self.recv_pipe = None
        self.send_pipe = None
        self.error_message = None
        # Guards the ``stopped`` / ``started_process`` pair. Without it, a
        # stop() landing while ``current_process.start()`` is still running
        # (slow on Windows, which spawns a fresh interpreter) sees
        # ``started_process`` still False, terminates nothing, and the
        # transcription runs to completion despite being canceled.
        self.process_lock = Lock()

    def transcribe(self) -> List[Segment]:
        time_started = datetime.datetime.now()
        logging.debug(
            "Starting whisper file transcription, task = %s", self.transcription_task
        )

        if torch.cuda.is_available():
            logging.debug(f"CUDA version detected: {torch.version.cuda}")

        self.recv_pipe, self.send_pipe = multiprocessing.Pipe(duplex=False)

        self.current_process = multiprocessing.Process(
            target=self.transcribe_whisper, args=(self.send_pipe, self.transcription_task)
        )
        with self.process_lock:
            canceled_before_start = self.stopped
            if not canceled_before_start:
                self.current_process.start()
                self.started_process = True

        if canceled_before_start:
            for pipe in (self.send_pipe, self.recv_pipe):
                try:
                    pipe.close()
                except OSError:
                    pass
            raise Exception("Transcription was canceled")

        self.read_line_thread = Thread(target=self.read_line, args=(self.recv_pipe,))
        self.read_line_thread.start()

        # Only join the process if it was actually started
        if self.started_process:
            self.current_process.join()

        # Close the send pipe after process ends to signal read_line thread to stop
        # This prevents the read thread from blocking on recv() after the process is gone
        try:
            if self.send_pipe and not self.send_pipe.closed:
                self.send_pipe.close()
        except OSError:
            pass

        # Close the receive pipe to unblock the read_line thread
        try:
            if self.recv_pipe and not self.recv_pipe.closed:
                self.recv_pipe.close()
        except OSError:
            pass

        # Join read_line_thread with timeout to prevent hanging
        if self.read_line_thread and self.read_line_thread.is_alive():
            self.read_line_thread.join(timeout=3)
            if self.read_line_thread.is_alive():
                logging.warning("Read line thread didn't terminate gracefully in transcribe()")

        self.started_process = False

        logging.debug(
            "whisper process completed with code = %s, time taken = %s,"
            " number of segments = %s",
            self.current_process.exitcode,
            datetime.datetime.now() - time_started,
            len(self.segments),
        )

        if self.current_process.exitcode != 0:
            # Check if the process was terminated (likely due to cancellation).
            # ``multiprocessing`` reports a process killed by a signal as the
            # negated signal number (-15 for SIGTERM, sent when the user stops
            # the transcription or closes the app), while shell-style exit codes
            # 124-143 report the same thing as 128 + signal number.
            if (
                self.current_process.exitcode < 0
                or self.current_process.exitcode in [124, 125, 126, 127, 128, 130, 137, 143]
            ):
                # Process was likely terminated, treat as cancellation
                logging.debug("Whisper process was terminated (exit code: %s), treating as cancellation", self.current_process.exitcode)
                raise Exception("Transcription was canceled")
            else:
                error = self.error_message or "Unknown error"
                logging.error("Whisper process failed (exit code: %s): %s", self.current_process.exitcode, error)
                raise Exception(error)

        return self.segments

    @classmethod
    def transcribe_whisper(
        cls, stderr_conn: Connection, task: FileTranscriptionTask
    ) -> None:
        # Patch subprocess on Windows to prevent console window flash
        # This is needed because multiprocessing spawns a new process without the main process patches
        if sys.platform == "win32":
            import subprocess
            _original_run = subprocess.run
            _original_popen = subprocess.Popen

            def _patched_run(*args, **kwargs):
                if 'startupinfo' not in kwargs:
                    si = subprocess.STARTUPINFO()
                    si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                    si.wShowWindow = subprocess.SW_HIDE
                    kwargs['startupinfo'] = si
                if 'creationflags' not in kwargs:
                    kwargs['creationflags'] = subprocess.CREATE_NO_WINDOW
                return _original_run(*args, **kwargs)

            class _PatchedPopen(subprocess.Popen):
                def __init__(self, *args, **kwargs):
                    if 'startupinfo' not in kwargs:
                        si = subprocess.STARTUPINFO()
                        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                        si.wShowWindow = subprocess.SW_HIDE
                        kwargs['startupinfo'] = si
                    if 'creationflags' not in kwargs:
                        kwargs['creationflags'] = subprocess.CREATE_NO_WINDOW
                    super().__init__(*args, **kwargs)

            subprocess.run = _patched_run
            subprocess.Popen = _PatchedPopen

        try:
            # Check if the file has audio streams before processing
            check_file_has_audio_stream(task.file_path)

            with pipe_stderr(stderr_conn):
                if task.transcription_options.model.model_type == ModelType.WHISPER_CPP:
                    segments = cls.transcribe_whisper_cpp(task)
                elif task.transcription_options.model.model_type == ModelType.HUGGING_FACE:
                    sys.stderr.write("0%\n")
                    segments = cls.transcribe_hugging_face(task)
                    sys.stderr.write("100%\n")
                elif (
                    task.transcription_options.model.model_type == ModelType.FASTER_WHISPER
                ):
                    segments = cls.transcribe_faster_whisper(task)
                elif task.transcription_options.model.model_type == ModelType.WHISPER:
                    segments = cls.transcribe_openai_whisper(task)
                else:
                    raise Exception(
                        f"Invalid model type: {task.transcription_options.model.model_type}"
                    )

                segments_json = json.dumps(segments, ensure_ascii=True, default=vars)
                sys.stderr.write(f"segments = {segments_json}\n")
                sys.stderr.write(WhisperFileTranscriber.READ_LINE_THREAD_STOP_TOKEN + "\n")
        except Exception as e:
            # Send error message back to the parent process
            stderr_conn.send(f"error = {str(e)}\n")
            stderr_conn.send(WhisperFileTranscriber.READ_LINE_THREAD_STOP_TOKEN + "\n")
            raise

    @classmethod
    def transcribe_whisper_cpp(cls, task: FileTranscriptionTask) -> List[Segment]:
        return WhisperCpp.transcribe(task)

    @classmethod
    def transcribe_hugging_face(cls, task: FileTranscriptionTask) -> List[Segment]:
        if not task.model_path:
            raise FileNotFoundError(
                "Hugging Face model is not available locally. "
                "The model download did not complete, try downloading it again."
            )

        model = TransformersTranscriber(task.model_path)

        # Handle language - MMS uses ISO 639-3 codes, Whisper uses ISO 639-1
        if model.is_mms_model:
            language = map_language_to_mms(task.transcription_options.language or "eng")
            # MMS only supports transcription, ignore translation task
            effective_task = Task.TRANSCRIBE.value
            # MMS doesn't support word-level timestamps
            word_timestamps = False
        else:
            language = (
                task.transcription_options.language
                if task.transcription_options.language is not None
                else "en"
            )
            effective_task = task.transcription_options.task.value
            word_timestamps = task.transcription_options.word_level_timings

        initial_prompt = "" if model.is_mms_model else (task.transcription_options.initial_prompt or "")
        result = model.transcribe(
            audio=task.file_path,
            language=language,
            task=effective_task,
            word_timestamps=word_timestamps,
            initial_prompt=initial_prompt,
        )
        return [
            Segment(
                start=int(segment.get("start") * 1000),
                end=int(segment.get("end") * 1000),
                text=segment.get("text"),
                translation=""
            )
            for segment in result.get("segments")
        ]

    @classmethod
    def transcribe_faster_whisper(cls, task: FileTranscriptionTask) -> List[Segment]:
        # Use the already-resolved local model path so we never hit the network
        model_size_or_path = task.model_path
        if not model_size_or_path:
            raise FileNotFoundError(
                "Faster Whisper model is not available locally. "
                "Check BUZZ_MODEL_ROOT and download the model into that cache first."
            )
            return []

        model_root_dir = user_cache_dir("Buzz")
        model_root_dir = os.path.join(model_root_dir, "models")
        model_root_dir = os.getenv("BUZZ_MODEL_ROOT", model_root_dir)
        force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
        if force_cpu != "false":
            os.environ["CUDA_VISIBLE_DEVICES"] = ""

        device = "auto"
        if not cuda_works():
            logging.debug("CUDA not available or not functional, using CPU")
            device = "cpu"
        elif torch.version.cuda < "12":
            logging.debug("Unsupported CUDA version (<12), using CPU")
            device = "cpu"

        if force_cpu != "false":
            device = "cpu"

        # Check if user wants reduced GPU memory usage (int8 quantization)
        reduce_gpu_memory = os.getenv("BUZZ_REDUCE_GPU_MEMORY", "false") != "false"
        compute_type = "default"
        if reduce_gpu_memory:
            compute_type = "int8" if device == "cpu" else "int8_float16"
            logging.debug(f"Using {compute_type} compute type for reduced memory usage")

        model = faster_whisper.WhisperModel(
            model_size_or_path=model_size_or_path,
            download_root=model_root_dir,
            device=device,
            compute_type=compute_type,
            cpu_threads=(os.cpu_count() or 8)//2,
        )

        audio = whisper_audio.load_audio(task.file_path)

        batched_model = faster_whisper.BatchedInferencePipeline(model=model)
        whisper_segments, info = batched_model.transcribe(
            audio=audio,
            language=task.transcription_options.language,
            task=task.transcription_options.task.value,
            # Prevent crash on Windows https://github.com/SYSTRAN/faster-whisper/issues/71#issuecomment-1526263764
            temperature = 0 if platform.system() == "Windows" else DEFAULT_WHISPER_TEMPERATURE,
            initial_prompt=task.transcription_options.initial_prompt,
            word_timestamps=task.transcription_options.word_level_timings,
            no_speech_threshold=0.4,
            log_progress=True,
        )
        segments = []
        for segment in whisper_segments:
            # Segment will contain words if word-level timings is True
            if segment.words:
                for word in segment.words:
                    segments.append(
                        Segment(
                            start=int(word.start * 1000),
                            end=int(word.end * 1000),
                            text=word.word,
                            translation=""
                        )
                    )
            else:
                segments.append(
                    Segment(
                        start=int(segment.start * 1000),
                        end=int(segment.end * 1000),
                        text=segment.text,
                        translation=""
                    )
                )

        return segments

    @classmethod
    def transcribe_openai_whisper(cls, task: FileTranscriptionTask) -> List[Segment]:
        force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
        if force_cpu != "false":
            os.environ["CUDA_VISIBLE_DEVICES"] = ""

        use_cuda = cuda_works() and force_cpu == "false"

        device = "cuda" if use_cuda else "cpu"

        # Monkeypatch torch.load to use weights_only=False for PyTorch 2.6+
        # This is required for loading Whisper models with the newer PyTorch versions
        original_torch_load = torch.load
        def patched_torch_load(*args, **kwargs):
            kwargs.setdefault('weights_only', False)
            return original_torch_load(*args, **kwargs)

        torch.load = patched_torch_load
        try:
            model = whisper.load_model(task.model_path, device=device)
        finally:
            torch.load = original_torch_load

        if task.transcription_options.word_level_timings:
            stable_whisper.modify_model(model)
            result: WhisperResult = model.transcribe(
                audio=whisper_audio.load_audio(task.file_path),
                language=task.transcription_options.language,
                task=task.transcription_options.task.value,
                temperature=DEFAULT_WHISPER_TEMPERATURE,
                initial_prompt=task.transcription_options.initial_prompt,
                no_speech_threshold=0.4,
                fp16=use_cuda,
            )
            return [
                Segment(
                    start=int(word.start * 1000),
                    end=int(word.end * 1000),
                    text=word.word.strip(),
                    translation=""
                )
                for segment in result.segments
                for word in segment.words
            ]

        result: dict = model.transcribe(
            audio=whisper_audio.load_audio(task.file_path),
            language=task.transcription_options.language,
            task=task.transcription_options.task.value,
            temperature=task.transcription_options.temperature,
            initial_prompt=task.transcription_options.initial_prompt,
            verbose=False,
            fp16=use_cuda,
        )
        segments = result.get("segments")
        return [
            Segment(
                start=int(segment.get("start") * 1000),
                end=int(segment.get("end") * 1000),
                text=segment.get("text"),
                translation=""
            )
            for segment in segments
        ]

    def stop(self):
        with self.process_lock:
            self.stopped = True
            # Read under the lock so a start() in progress has finished (and
            # published the pid) before we decide there is nothing to kill.
            started_process = self.started_process

        if started_process:
            # Kill the whisper-cli subprocess the worker spawned first. The
            # worker's own terminate() below does not reach it, so it would
            # otherwise keep running (orphaned) after Stop / app close.
            if self.current_process.pid is not None:
                terminate_child_processes(self.current_process.pid)

            # Terminate the worker itself; it is reaped via join() below.
            self.current_process.terminate()

            # Close the pipes to unblock the read_line thread, which is
            # otherwise stuck in recv() until it gets EOF / an OSError.
            try:
                if hasattr(self, 'send_pipe') and self.send_pipe:
                    self.send_pipe.close()
            except Exception as e:
                logging.debug(f"Error closing send_pipe: {e}")

            try:
                if hasattr(self, 'recv_pipe') and self.recv_pipe:
                    self.recv_pipe.close()
            except Exception as e:
                logging.debug(f"Error closing recv_pipe: {e}")

            if self.read_line_thread and self.read_line_thread.is_alive():
                self.read_line_thread.join(timeout=5)
                if self.read_line_thread.is_alive():
                    logging.warning("Read line thread still alive after 5s")

            self.current_process.join(timeout=10)
            if self.current_process.is_alive():
                logging.warning("Process didn't terminate gracefully, force killing")
                self.current_process.kill()
                self.current_process.join(timeout=5)

    def read_line(self, pipe: Connection):
        while True:
            try:
                line = pipe.recv().strip()

                # Uncomment to debug
                # print(f"*** DEBUG ***: {line}")

            except (EOFError, BrokenPipeError, ConnectionResetError, OSError):
                # Connection closed, broken, or process crashed (Windows RPC errors raise OSError)
                break
            except Exception as e:
                logging.debug(f"Error reading from pipe: {e}")
                break

            if line == self.READ_LINE_THREAD_STOP_TOKEN:
                return

            if line.startswith("segments = "):
                segments_dict = json.loads(line[11:])
                segments = [
                    Segment(
                        start=segment.get("start"),
                        end=segment.get("end"),
                        text=segment.get("text"),
                        translation=""
                    )
                    for segment in segments_dict
                ]
                self.segments = segments
            elif line.startswith("error = "):
                self.error_message = line[8:]
            else:
                try:
                    match = PROGRESS_REGEX.search(line)
                    if match is not None:
                        progress = int(match.group().strip("%"))
                        self.progress.emit((progress, 100))
                except ValueError:
                    logging.debug("whisper (stderr): %s", line)
                    continue