1+ import asyncio
2+ import websockets
3+ import json
4+ import base64
5+ import wave
6+ import os
7+
8+ def wav_to_base64 (wav_path ):
9+ """Convert a WAV file to base64 encoded string"""
10+ with open (wav_path , "rb" ) as wav_file :
11+ return base64 .b64encode (wav_file .read ()).decode ('utf-8' )
12+
13+ async def send_websocket_data ():
14+ # Connection details
15+ uri = "wss://model-rwn1jgd3.api.baseten.co/v1/websocket"
16+ headers = {"Authorization" : "Api-Key vVolDAU0.Mbynm8M7VGnaGqLbW9pwfWxFePNrGw8G" }
17+
18+ async with websockets .connect (uri , extra_headers = headers ) as websocket :
19+ # For the TTS model, we send text instead of audio
20+ text_data = {
21+ "text" : "Hello, this is a test of the text to speech websocket API." ,
22+ "language" : "en" ,
23+ "chunk_size" : 20
24+ }
25+
26+ # Send the text data as JSON
27+ await websocket .send (json .dumps (text_data ))
28+ print (f"Sent text: { text_data ['text' ]} )
29+
30+ # Collect audio chunks
31+ audio_chunks = []
32+
33+ # Process responses
34+ while True :
35+ try :
36+ response = await websocket .recv ()
37+
38+ # Try to parse as JSON
39+ try :
40+ data = json .loads (response )
41+ print (f"Received response: { data .get ('type' , 'unknown' )} )
42+
43+ if data .get ("type" ) == "chunk" :
44+ # Decode and save the audio chunk
45+ audio_chunk = base64 .b64decode (data ["data" ])
46+ audio_chunks .append (audio_chunk )
47+ print ("Saved audio chunk" )
48+
49+ elif data .get ("type" ) == "complete" :
50+ print (f"Processing complete: { data .get ('message' )} )
51+ break
52+
53+ elif data .get ("type" ) == "error" :
54+ print (f"Error: { data .get ('message' )} )
55+ break
56+
57+ except json .JSONDecodeError :
58+ # Not JSON, print the first part
59+ print (f"Received non-JSON response: { response [:50 ]} )
60+ break
61+
62+ except Exception as e :
63+ print (f"Error receiving data: { str (e )} )
64+ break
65+
66+ # Save the audio to a WAV file if we received chunks
67+ if audio_chunks :
68+ output_file = "tts_output.wav"
69+ with wave .open (output_file , 'wb' ) as wf :
70+ wf .setnchannels (1 ) # Mono
71+ wf .setsampwidth (2 ) # 16-bit
72+ wf .setframerate (24000 ) # XTTS default sample rate
73+ wf .writeframes (b'' .join (audio_chunks ))
74+
75+ print (f"Audio saved to { output_file } )
76+ print (f"Full path: { os .path .abspath (output_file )} )
77+ else :
78+ print ("No audio data received" )
79+
80+ async def test_multiple_concurrent_requests ():
81+ """Test sending multiple concurrent requests to the TTS websocket API"""
82+
83+ async def single_request (idx ):
84+ """Handle a single request with unique text and output file"""
85+ output_file = f"tts_output_{ idx }
86+ text = f"This is concurrent test number { idx }
87+
88+ try :
89+ # Connection details
90+ uri = "wss://model-rwn1jgd3.api.baseten.co/v1/websocket"
91+ headers = {"Authorization" : "Api-Key vVolDAU0.Mbynm8M7VGnaGqLbW9pwfWxFePNrGw8G" }
92+
93+ async with websockets .connect (uri , extra_headers = headers ) as websocket :
94+ # Send text data as JSON
95+ text_data = {
96+ "text" : text ,
97+ "language" : "en" ,
98+ "chunk_size" : 20
99+ }
100+
101+ await websocket .send (json .dumps (text_data ))
102+ print (f"Request { idx } { text } )
103+
104+ # Collect audio chunks
105+ audio_chunks = []
106+
107+ # Process responses
108+ while True :
109+ try :
110+ response = await websocket .recv ()
111+
112+ # Try to parse as JSON
113+ try :
114+ data = json .loads (response )
115+
116+ if data .get ("type" ) == "chunk" :
117+ # Decode and save the audio chunk
118+ audio_chunk = base64 .b64decode (data ["data" ])
119+ audio_chunks .append (audio_chunk )
120+
121+ elif data .get ("type" ) == "complete" :
122+ print (f"Request { idx } )
123+ break
124+
125+ elif data .get ("type" ) == "error" :
126+ print (f"Request { idx } { data .get ('message' )} )
127+ return False
128+
129+ except json .JSONDecodeError :
130+ print (f"Request { idx } )
131+ return False
132+
133+ except Exception as e :
134+ print (f"Request { idx } { str (e )} )
135+ return False
136+
137+ # Save the audio to a WAV file if we received chunks
138+ if audio_chunks :
139+ with wave .open (output_file , 'wb' ) as wf :
140+ wf .setnchannels (1 ) # Mono
141+ wf .setsampwidth (2 ) # 16-bit
142+ wf .setframerate (24000 ) # XTTS default sample rate
143+ wf .writeframes (b'' .join (audio_chunks ))
144+
145+ print (f"Request { idx } { output_file } )
146+ return True
147+ else :
148+ print (f"Request { idx } )
149+ return False
150+
151+ except Exception as e :
152+ print (f"Request { idx } { str (e )} )
153+ return False
154+
155+ num_requests = 4
156+
157+ print (f"Starting { num_requests } )
158+ results = await asyncio .gather (* [single_request (i + 1 ) for i in range (num_requests )])
159+
160+ successful = results .count (True )
161+ print (f"Completed { successful } { num_requests } )
162+ return successful == num_requests
163+
164+ # Run the tests
165+ if __name__ == "__main__" :
166+ asyncio .run (send_websocket_data ())
167+ print ("\n --- Testing multiple concurrent requests ---\n " )
168+ asyncio .run (test_multiple_concurrent_requests ())
0 commit comments