piealamodewhitebread/SillyTavern-Extras1

1"""2Speech-to-text module based on Vosk and Whisper for SillyTavern Extras3    - Vosk website: https://alphacephei.com/vosk/4    - Vosk api: https://github.com/alphacep/vosk-api5    - Whisper github: https://github.com/openai/whisper6 7Authors:8    - Tony Ribeiro (https://github.com/Tony-sama)9 10Models are saved into user cache folder, example: C:/Users/toto/.cache/whisper and C:/Users/toto/.cache/vosk11 12References:13    - Code adapted from:14        - whisper github: https://github.com/openai/whisper15        - oobabooga text-generation-webui github: https://github.com/oobabooga/text-generation-webui16        - vosk github: https://github.com/alphacep/vosk-api/blob/master/python/example/test_microphone.py17"""18from flask import jsonify, abort19 20import queue21import sys22import sounddevice as sd23import soundfile as sf24import io25import numpy as np26from scipy.io.wavfile import write27 28import vosk29import whisper30 31DEBUG_PREFIX = "<stt streaming module>"32RECORDING_FILE_PATH = "stt_test.wav"33 34whisper_model = None35vosk_model = None36device = None37 38def load_model(file_path=None):39    """40    Load given vosk model from file or default to en-us model.41    Download model to user cache folder, example: C:/Users/toto/.cache/vosk42    """43 44    if file_path is None:45        return (whisper.load_model("base.en"), vosk.Model(lang="en-us"))46    else:47        return (whisper.load_model(file_path), vosk.Model(lang="en-us"))48    49def convert_bytearray_to_wav_ndarray(input_bytearray: bytes, sampling_rate=16000):50    """51    Convert a bytearray to wav format to output in a file for quality check debuging52    """53    bytes_wav = bytes()54    byte_io = io.BytesIO(bytes_wav)55    write(byte_io, sampling_rate, np.frombuffer(input_bytearray, dtype=np.int16))56    output_wav = byte_io.read()57    output, _ = sf.read(io.BytesIO(output_wav))58    return output59 60def record_and_transcript():61    """62    Continuously record from mic and transcript voice.63    Return the transcript once no more voice is detected.64    """65    if whisper_model is None:66        print(DEBUG_PREFIX,"Whisper model not initialized yet.")67        return ""68    69    q = queue.Queue()70    stream_errors = list()71 72    def callback(indata, frames, time, status):73        """This is called (from a separate thread) for each audio block."""74        if status:75            print(status, file=sys.stderr)76            stream_errors.append(status)77        q.put(bytes(indata))78 79    try:80        device_info = sd.query_devices(device, "input")81        # soundfile expects an int, sounddevice provides a float:82        samplerate = int(device_info["default_samplerate"])83 84        print(DEBUG_PREFIX, "Start recording from:", device_info["name"], "with samplerate", samplerate)85 86        with sd.RawInputStream(samplerate=samplerate, blocksize = 8000, device=device, dtype="int16", channels=1, callback=callback):87 88            rec = vosk.KaldiRecognizer(vosk_model, samplerate)89            full_recording = bytearray()90            while True:91                data = q.get()92                if len(stream_errors) > 0:93                    raise Exception(DEBUG_PREFIX+" Stream errors: "+str(stream_errors))94                95                full_recording.extend(data)96 97                if rec.AcceptWaveform(data):98                    # Extract transcript string99                    transcript = rec.Result()[14:-3]100                    print(DEBUG_PREFIX, "Transcripted from microphone stream (vosk):", transcript)101 102                    # ----------------------------------103                    # DEBUG: save recording to wav file104                    # ----------------------------------105                    output_file = convert_bytearray_to_wav_ndarray(input_bytearray=full_recording, sampling_rate=samplerate)106                    sf.write(file=RECORDING_FILE_PATH, data=output_file, samplerate=samplerate)107                    print(DEBUG_PREFIX, "Recorded message saved to", RECORDING_FILE_PATH)108                    109                    # Whisper HACK110                    result = whisper_model.transcribe(RECORDING_FILE_PATH)111                    transcript = result["text"]112                    print(DEBUG_PREFIX, "Transcripted from audio file (whisper):", transcript)113                    # ----------------------------------114 115                    return jsonify({"transcript": transcript})116                #else:117                #    print(rec.PartialResult())118 119    except Exception as e: # No exception observed during test but we never know120        print(e)121        abort(500, DEBUG_PREFIX+" Exception occurs while recording")