Spaces:
Runtime error
Runtime error
File size: 4,662 Bytes
b5b0b9a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import threading
from queue import Queue, Empty
import numpy as np
import requests
import base64
import time
from dataclasses import dataclass, field
import websocket
import threading
import ssl
import librosa
import os
class AudioStreamingClient:
def __init__(self):
self.auth_token = os.environ.get("HF_AUTH_TOKEN", None)
self.api_url = os.environ.get("HF_API_URL", None)
self.stop_event = threading.Event()
self.send_queue = Queue()
self.recv_queue = Queue()
self.session_id = None
self.headers = {
"Accept": "application/json",
"Authorization": f"Bearer {self.auth_token}",
"Content-Type": "application/json"
}
self.session_state = "idle" # Possible states: idle, sending, processing, waiting
self.ws_ready = threading.Event()
def start(self):
print("Starting audio streaming...")
ws_url = self.api_url.replace("http", "ws") + "/ws"
self.ws = websocket.WebSocketApp(
ws_url,
header=[f"{key}: {value}" for key, value in self.headers.items()],
on_open=self.on_open,
on_message=self.on_message,
on_error=self.on_error,
on_close=self.on_close
)
self.ws_thread = threading.Thread(target=self.ws.run_forever, kwargs={'sslopt': {"cert_reqs": ssl.CERT_NONE}})
self.ws_thread.start()
# Wait for the WebSocket to be ready
self.ws_ready.wait()
self.send_thread = threading.Thread(target=self.send_audio)
self.send_thread.start()
def on_close(self):
self.stop_event.set()
self.send_thread.join()
self.ws.close()
self.ws_thread.join()
print("Audio streaming stopped.")
def on_open(self, ws):
print("WebSocket connection opened.")
self.ws_ready.set() # Signal that the WebSocket is ready
def on_message(self, ws, message):
# message is bytes
if message == b'DONE':
print("listen")
self.session_state = "listen"
else:
print("processing")
self.session_state = "processing"
audio_np = np.frombuffer(message, dtype=np.int16)
self.recv_queue.put(audio_np)
def on_error(self, ws, error):
print(f"WebSocket error: {error}")
def on_close(self, ws, close_status_code, close_msg):
print("WebSocket connection closed.")
def send_audio(self):
while not self.stop_event.is_set():
if not self.send_queue.empty():
chunk = self.send_queue.get()
if self.session_state != "processing":
self.ws.send(chunk.tobytes(), opcode=websocket.ABNF.OPCODE_BINARY)
else:
self.ws.send([], opcode=websocket.ABNF.OPCODE_BINARY) # handshake
time.sleep(0.01)
def put_audio(self, chunk, sample_rate):
chunk = np.clip(chunk, -32768, 32767).astype(np.int16)
chunk = chunk.astype(np.float32) / 32768.0
chunk = librosa.resample(chunk, orig_sr=48000, target_sr=16000)
chunk = (chunk * 32768.0).astype(np.int16)
self.send_queue.put(chunk)
def get_audio(self, sample_rate, output_size):
output_chunk = np.array([], dtype=np.int16)
output_sample_rate = 16000
output_chunk_size = int(output_size*output_sample_rate/sample_rate)
while output_chunk.size < output_chunk_size:
try:
self.ws.send([], opcode=websocket.ABNF.OPCODE_BINARY) # handshake
chunk = self.recv_queue.get(timeout=0.1)
except Empty:
chunk = None
if chunk is not None:
# Ensure chunk is int16 and clip to valid range
chunk_int16 = np.clip(chunk, -32768, 32767).astype(np.int16)
output_chunk = np.concatenate([output_chunk, chunk_int16])
else:
print("padding chunk of size ", len(output_chunk))
output_chunk = np.pad(output_chunk, (0, output_chunk_size - len(output_chunk)))
output_chunk = output_chunk.astype(np.float32) / 32768.0
output_chunk = librosa.resample(output_chunk, orig_sr=output_sample_rate, target_sr=sample_rate)
output_chunk = (output_chunk * 32768.0).astype(np.int16)
print("output_chunk size: ", len(output_chunk))
output_chunk = output_chunk[:output_size]
return np.pad(output_chunk, (0, output_size - len(output_chunk)))
if __name__ == "__main__":
client = AudioStreamingClient()
client.start() |