WIP

faster-whisper-server-client.py

@@ -0,0 +1,225 @@
+import argparse
+import json
+import threading
+import time
+import websocket
+import os
+
+import librosa
+import numpy as np
+
+# Define the default WebSocket endpoint
+DEFAULT_WS_URL = "ws://localhost:8000/v1/audio/transcriptions"
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Stream audio to the transcription WebSocket endpoint.")
+    parser.add_argument("audio_file", help="Path to the input audio file.")
+    parser.add_argument("--url", default=DEFAULT_WS_URL, help="WebSocket endpoint URL.")
+    parser.add_argument("--model", type=str, help="Model name to use for transcription.")
+    parser.add_argument("--language", type=str, help="Language code for transcription.")
+    parser.add_argument(
+        "--response_format",
+        type=str,
+        default="verbose_json",
+        choices=["text", "json", "verbose_json"],
+        help="Response format.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.0, help="Temperature for transcription.")
+    parser.add_argument("--vad_filter", action="store_true", help="Enable voice activity detection filter.")
+    parser.add_argument("--chunk_duration", type=float, default=1.0, help="Duration of each audio chunk in seconds.")
+    return parser.parse_args()
+
+
+def preprocess_audio(audio_file, target_sr=16000):
+    """
+    Load the audio file, convert to mono 16kHz, and return the audio data.
+    """
+    if audio_file.endswith(".mp3"):
+        # Convert MP3 to WAV using ffmpeg
+        wav_file = audio_file.replace(".mp3", ".wav")
+        if not os.path.exists(wav_file):
+            command = f'ffmpeg -i "{audio_file}" -ac 1 -ar {target_sr} "{wav_file}"'
+            print(f"Converting MP3 to WAV: {command}")
+            os.system(command)
+        audio_file = wav_file
+
+    print(f"Loading audio file {audio_file}")
+    audio_data, sr = librosa.load(audio_file, sr=target_sr, mono=True)
+    return audio_data, sr
+
+def chunk_audio(audio_data, sr, chunk_duration):
+    """
+    Split the audio data into chunks of specified duration.
+    """
+    chunk_samples = int(chunk_duration * sr)
+    total_samples = len(audio_data)
+    chunks = [
+        audio_data[i:i + chunk_samples]
+        for i in range(0, total_samples, chunk_samples)
+    ]
+    print(f"Split audio into {len(chunks)} chunks of {chunk_duration} seconds each.")
+    return chunks
+
+
+def build_query_params(args):
+    """
+    Build the query parameters for the WebSocket URL based on command-line arguments.
+    """
+    params = {}
+    if args.model:
+        params["model"] = args.model
+    if args.language:
+        params["language"] = args.language
+    if args.response_format:
+        params["response_format"] = args.response_format
+    if args.temperature is not None:
+        params["temperature"] = str(args.temperature)
+    if args.vad_filter:
+        params["vad_filter"] = "true"
+    return params
+
+
+def websocket_url_with_params(base_url, params):
+    """
+    Append query parameters to the WebSocket URL.
+    """
+    from urllib.parse import urlencode
+
+    if params:
+        query_string = urlencode(params)
+        url = f"{base_url}?{query_string}"
+    else:
+        url = base_url
+    return url
+
+
+def on_message(ws, message):
+    """
+    Callback function when a message is received from the server.
+    """
+    try:
+        data = json.loads(message)
+        # Accumulate transcriptions
+        if ws.args.response_format == "verbose_json":
+            segments = data.get('segments', [])
+            ws.transcriptions.extend(segments)
+            for segment in segments:
+                print(f"Received segment: {segment['text']}")
+        else:
+            # For 'json' or 'text' format
+            ws.transcriptions.append(data)
+            print(f"Transcription: {data['text']}")
+    except json.JSONDecodeError:
+        print(f"Received non-JSON message: {message}")
+
+
+def on_error(ws, error):
+    """
+    Callback function when an error occurs.
+    """
+    print(f"WebSocket error: {error}")
+
+
+def on_close(ws, close_status_code, close_msg):
+    """
+    Callback function when the WebSocket connection is closed.
+    """
+    print("WebSocket connection closed")
+
+
+def on_open(ws):
+    """
+    Callback function when the WebSocket connection is opened.
+    """
+    print("WebSocket connection opened")
+    ws.transcriptions = []  # Initialize the list to store transcriptions
+
+
+def send_audio_chunks(ws, audio_chunks, sr):
+    """
+    Send audio chunks to the WebSocket server.
+    """
+    for idx, chunk in enumerate(audio_chunks):
+        # Ensure little-endian format
+        audio_bytes = chunk.astype('<f4').tobytes()
+        ws.send(audio_bytes, opcode=websocket.ABNF.OPCODE_BINARY)
+        print(f"Sent chunk {idx + 1}/{len(audio_chunks)}")
+        time.sleep(0.1)  # Small delay to simulate real-time streaming
+    print("All audio chunks sent")
+    # Optionally, wait to receive remaining messages
+    time.sleep(2)
+    ws.close()
+    print("Closed WebSocket connection")
+
+
+
+def format_timestamp(seconds):
+    """
+    Convert seconds to SRT timestamp format (HH:MM:SS,mmm).
+    """
+    total_milliseconds = int(seconds * 1000)
+    hours = total_milliseconds // (3600 * 1000)
+    minutes = (total_milliseconds % (3600 * 1000)) // (60 * 1000)
+    secs = (total_milliseconds % (60 * 1000)) // 1000
+    milliseconds = total_milliseconds % 1000
+    return f"{hours:02}:{minutes:02}:{secs:02},{milliseconds:03}"
+
+
+def generate_srt(transcriptions):
+    """
+    Generate and print SRT content from transcriptions.
+    """
+    print("\nGenerated SRT:")
+    for idx, segment in enumerate(transcriptions, 1):
+        start_time = format_timestamp(segment['start'])
+        end_time = format_timestamp(segment['end'])
+        text = segment['text']
+        print(f"{idx}")
+        print(f"{start_time} --> {end_time}")
+        print(f"{text}\n")
+
+
+def run_websocket_client(args):
+    """
+    Run the WebSocket client to stream audio and receive transcriptions.
+    """
+    audio_data, sr = preprocess_audio(args.audio_file)
+    audio_chunks = chunk_audio(audio_data, sr, args.chunk_duration)
+
+    params = build_query_params(args)
+    ws_url = websocket_url_with_params(args.url, params)
+
+    ws = websocket.WebSocketApp(
+        ws_url,
+        on_open=on_open,
+        on_message=on_message,
+        on_error=on_error,
+        on_close=on_close,
+    )
+    ws.args = args  # Attach args to ws to access inside callbacks
+
+    # Run the WebSocket in a separate thread to allow sending and receiving simultaneously
+    ws_thread = threading.Thread(target=ws.run_forever)
+    ws_thread.start()
+
+    # Wait for the connection to open
+    while not ws.sock or not ws.sock.connected:
+        time.sleep(0.1)
+
+    # Send the audio chunks
+    send_audio_chunks(ws, audio_chunks, sr)
+
+    # Wait for the WebSocket thread to finish
+    ws_thread.join()
+
+    # Generate SRT if transcriptions are available
+    if hasattr(ws, 'transcriptions') and ws.transcriptions:
+        generate_srt(ws.transcriptions)
+    else:
+        print("No transcriptions received.")
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    run_websocket_client(args)

ws_server.py

@@ -0,0 +1,175 @@
+# Import the necessary components from whisper_online.py
+import logging
+import os
+
+import librosa
+import soundfile
+import uvicorn
+from fastapi import FastAPI, WebSocket
+from starlette.websockets import WebSocketDisconnect
+
+from libs.whisper_streaming.whisper_online import (
+    ASRBase,
+    OnlineASRProcessor,
+    VACOnlineASRProcessor,
+    add_shared_args,
+    asr_factory,
+    set_logging,
+    create_tokenizer,
+    load_audio,
+    load_audio_chunk, OpenaiApiASR,
+    set_logging
+)
+
+import argparse
+import sys
+import numpy as np
+import io
+import soundfile as sf
+import wave
+import requests
+import argparse
+
+logger = logging.getLogger(__name__)
+
+SAMPLING_RATE = 16000
+WARMUP_FILE = "mono16k.test_hebrew.wav"
+AUDIO_FILE_URL = "https://raw.githubusercontent.com/AshDavid12/runpod-serverless-forked/main/test_hebrew.wav"
+
+is_first = True
+asr, online = None, None
+min_limit = None  # min_chunk*SAMPLING_RATE
+app = FastAPI()
+
+
+def convert_to_mono_16k(input_wav: str, output_wav: str) -> None:
+    """
+    Converts any .wav file to mono 16 kHz.
+
+    Args:
+        input_wav (str): Path to the input .wav file.
+        output_wav (str): Path to save the output .wav file with mono 16 kHz.
+    """
+    # Step 1: Load the audio file with librosa
+    audio_data, original_sr = librosa.load(input_wav, sr=None, mono=False)  # Load at original sampling rate
+    logger.info("Loaded audio with shape: %s, original sampling rate: %d" % (audio_data.shape, original_sr))
+
+    # Step 2: If the audio has multiple channels, average them to make it mono
+    if audio_data.ndim > 1:
+        audio_data = librosa.to_mono(audio_data)
+
+    # Step 3: Resample the audio to 16 kHz
+    resampled_audio = librosa.resample(audio_data, orig_sr=original_sr, target_sr=SAMPLING_RATE)
+
+    # Step 4: Save the resampled audio as a .wav file in mono at 16 kHz
+    sf.write(output_wav, resampled_audio, SAMPLING_RATE)
+
+    logger.info(f"Converted audio saved to {output_wav}")
+
+def download_warmup_file():
+    # Download the audio file if not already present
+    audio_file_path = "test_hebrew.wav"
+    if not os.path.exists(WARMUP_FILE):
+        if not os.path.exists(audio_file_path):
+            response = requests.get(AUDIO_FILE_URL)
+            with open(audio_file_path, 'wb') as f:
+                f.write(response.content)
+
+        convert_to_mono_16k(audio_file_path, WARMUP_FILE)
+
+
+async def receive_audio_chunk(self, websocket: WebSocket):
+    # receive all audio that is available by this time
+    # blocks operation if less than self.min_chunk seconds is available
+    # unblocks if connection is closed or a chunk is available
+    out = []
+    while sum(len(x) for x in out) < min_limit:
+        raw_bytes = await websocket.receive_bytes()
+        if not raw_bytes:
+            break
+
+        sf = soundfile.SoundFile(io.BytesIO(raw_bytes), channels=1,endian="LITTLE",samplerate=SAMPLING_RATE, subtype="PCM_16",format="RAW")
+        audio, _ = librosa.load(sf,sr=SAMPLING_RATE,dtype=np.float32)
+        out.append(audio)
+
+    if not out:
+        return None
+
+    conc = np.concatenate(out)
+    if self.is_first and len(conc) < min_limit:
+        return None
+    self.is_first = False
+    return conc
+
+# Define WebSocket endpoint
+@app.websocket("/ws_transcribe_streaming")
+async def websocket_transcribe(websocket: WebSocket):
+    logger.info("New WebSocket connection request received.")
+    await websocket.accept()
+    logger.info("WebSocket connection established successfully.")
+
+    asr, online = asr_factory(args)
+
+    # warm up the ASR because the very first transcribe takes more time than the others.
+    # Test results in https://github.com/ufal/whisper_streaming/pull/81
+    a = load_audio_chunk(WARMUP_FILE, 0, 1)
+    asr.transcribe(a)
+    logger.info("Whisper is warmed up.")
+    global min_limit
+    min_limit = args.min_chunk_size * SAMPLING_RATE
+
+    try:
+        out = []
+        while True:
+            try:
+                # Receive JSON data
+                raw_bytes = await websocket.receive_json()
+
+                sf = soundfile.SoundFile(io.BytesIO(raw_bytes), channels=1, endian="LITTLE", samplerate=SAMPLING_RATE,
+                                         subtype="PCM_16", format="RAW")
+                audio, _ = librosa.load(sf, sr=SAMPLING_RATE, dtype=np.float32)
+                out.append(audio)
+
+                # Call the transcribe function
+                # segments, info = await asyncio.to_thread(model.transcribe,
+                segments, info = model.transcribe(
+                    audio_file_path,
+                    language='he',
+                    initial_prompt=input_data.init_prompt,
+                    beam_size=5,
+                    word_timestamps=True,
+                    condition_on_previous_text=True
+                )
+
+                # Convert segments to list and serialize
+                segments_list = list(segments)
+                segments_serializable = [segment_to_dict(s) for s in segments_list]
+                logger.info(get_raw_words_from_segments(segments_list))
+                # Send the serialized segments back to the client
+                await websocket.send_json(segments_serializable)
+
+            except WebSocketDisconnect:
+                logger.info("WebSocket connection closed by the client.")
+                break
+            except Exception as e:
+                logger.error(f"Unexpected error during WebSocket transcription: {e}")
+                await websocket.send_json({"error": str(e)})
+    finally:
+        logger.info("Cleaning up and closing WebSocket connection.")
+
+def main():
+    args = argparse.ArgumentParser()
+    args = add_shared_args(args)
+    args.parse_args([
+        '--lan', 'he',
+        '--model', 'ivrit-ai/faster-whisper-v2-d4',
+        '--backend', 'faster-whisper',
+        '--vad',
+        # '--vac', '--buffer_trimming', 'segment', '--buffer_trimming_sec', '15', '--min_chunk_size', '1.0', '--vac_chunk_size', '0.04', '--start_at', '0.0', '--offline', '--comp_unaware', '--log_level', 'DEBUG'
+    ])
+
+
+    global asr, online
+
+
+    uvicorn.run(app)