gajim-plugins/stt_voice_messages/model.py

# This file is part of Gajim.
#
# Gajim is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Gajim is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Gajim. If not, see <http://www.gnu.org/licenses/>.

from __future__ import annotations

from typing import Any

import gc
import logging
import threading
from collections.abc import Callable
from importlib.util import find_spec
from pathlib import Path

import numpy as np
from gi.repository import GLib
from gi.repository import Gst

log = logging.getLogger("gajim.p.sttvm_model")

_IDLE_UNLOAD_SECONDS = 300


def load_audio(path: Path, sample_rate: int = 16000) -> np.ndarray:
    Gst.init(None)
    pipeline = Gst.parse_launch(
        "filesrc name=src ! decodebin ! audioconvert ! audioresample ! "
        f"audio/x-raw,format=F32LE,rate={sample_rate},channels=1 ! "
        "appsink name=sink sync=false"
    )

    pipeline.get_by_name("src").set_property("location", str(path))
    sink = pipeline.get_by_name("sink")

    chunks: list[np.ndarray] = []
    try:
        pipeline.set_state(Gst.State.PLAYING)
        while (sample := sink.emit("try-pull-sample", 10 * Gst.SECOND)) is not None:
            buf = sample.get_buffer()
            _, info = buf.map(Gst.MapFlags.READ)
            chunks.append(np.frombuffer(bytes(info.data), dtype=np.float32))
            buf.unmap(info)
        pipeline.set_state(Gst.State.NULL)
        # wait for async cleanup so decodebin/appsink release memory
        pipeline.get_state(Gst.CLOCK_TIME_NONE)
    finally:
        del pipeline
        gc.collect()

    if not chunks:
        raise RuntimeError(f"Could not decode audio: {path}")
    return np.concatenate(chunks)


class OnnxAsrModel:
    def __init__(self, model_id: str, quantization: str) -> None:
        self._model_id = model_id
        self._quantization = quantization
        self._model: Any = None
        self._loaded = False
        self._busy = False
        self._recognize_lock = threading.Lock()
        self._unload_source: int | None = None

    @staticmethod
    def available() -> bool:
        return find_spec("onnx_asr") is not None

    @property
    def is_loaded(self) -> bool:
        return self._loaded

    @property
    def model_id(self) -> str:
        return self._model_id

    @property
    def will_download(self) -> bool:
        if self._loaded or not self.available():
            return False
        from huggingface_hub import try_to_load_from_cache
        from onnx_asr.resolver import model_repos

        repo = model_repos.get(self._model_id, self._model_id)
        if "/" not in repo:
            # Local path, nothing to download
            return False
        return not isinstance(try_to_load_from_cache(repo, "config.json"), str)

    def set_config(self, model_id: str, quantization: str) -> None:
        if model_id == self._model_id and quantization == self._quantization:
            return
        self.unload_now()
        self._model_id = model_id
        self._quantization = quantization

    def load(self) -> None:
        if self._loaded:
            self._schedule_unload()
            return

        import onnx_asr

        log.debug("Loading model %s", self._model_id)
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        quantization = None if self._quantization == "fp32" else self._quantization
        model = onnx_asr.load_model(self._model_id, quantization=quantization, providers=providers)
        self._model = model.with_vad(onnx_asr.load_vad("silero"))
        self._loaded = True
        self._schedule_unload()

    def recognize(
        self,
        audio_file: Path,
        on_partial: Callable[[str], None] | None = None,
    ) -> str:
        # global lock, one transcription at a time
        with self._recognize_lock:
            self.load()
            self._busy = True
            try:
                audio = load_audio(audio_file)
                parts: list[str] = []
                text = ""
                for segment in self._model.recognize(audio, sample_rate=16000):
                    parts.append(segment.text.strip())
                    text = " ".join(part for part in parts if part)
                    if on_partial is not None:
                        on_partial(text)
            finally:
                self._busy = False
        self._schedule_unload()
        return text

    def unload_now(self) -> None:
        self._model = None
        self._loaded = False
        gc.collect()

    def _schedule_unload(self) -> None:
        if self._unload_source is not None:
            GLib.source_remove(self._unload_source)
        self._unload_source = GLib.timeout_add_seconds(
            _IDLE_UNLOAD_SECONDS, self._unload_idle
        )

    def _unload_idle(self) -> bool:
        if self._busy:
            return GLib.SOURCE_CONTINUE
        self._unload_source = None
        self.unload_now()
        return GLib.SOURCE_REMOVE