@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 182 lines 5.42 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API.
4
5 Start the vLLM server first:
6
7 vllm serve mistralai/Voxtral-Mini-4B-Realtime-2602 --enforce-eager
8
9 Then run this script:
10
11 python openai_realtime_microphone_client.py --host localhost --port 8000
12
13 Use --share to create a public Gradio link.
14
15 Requirements: websockets, numpy, gradio
16 """
17
18 import argparse
19 import asyncio
20 import json
21 import queue
22 import threading
23
24 import gradio as gr
25 import numpy as np
26 import pybase64 as base64
27 import websockets
28
29 SAMPLE_RATE = 16_000
30
31 # Global state
32 audio_queue: queue.Queue = queue.Queue()
33 transcription_text = ""
34 is_running = False
35 ws_url = ""
36 model = ""
37
38
39 async def websocket_handler():
40 """Connect to WebSocket and handle audio streaming + transcription."""
41 global transcription_text, is_running
42
43 async with websockets.connect(ws_url) as ws:
44 # Wait for session.created
45 await ws.recv()
46
47 # Validate model
48 await ws.send(json.dumps({"type": "session.update", "model": model}))
49
50 # Signal ready
51 await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
52
53 async def send_audio():
54 while is_running:
55 try:
56 chunk = await asyncio.get_event_loop().run_in_executor(
57 None, lambda: audio_queue.get(timeout=0.1)
58 )
59 await ws.send(
60 json.dumps(
61 {"type": "input_audio_buffer.append", "audio": chunk}
62 )
63 )
64 except queue.Empty:
65 continue
66
67 async def receive_transcription():
68 global transcription_text
69 async for message in ws:
70 data = json.loads(message)
71 if data.get("type") == "transcription.delta":
72 transcription_text += data["delta"]
73
74 await asyncio.gather(send_audio(), receive_transcription())
75
76
77 def start_websocket():
78 """Start WebSocket connection in background thread."""
79 global is_running
80 is_running = True
81 loop = asyncio.new_event_loop()
82 asyncio.set_event_loop(loop)
83 try:
84 loop.run_until_complete(websocket_handler())
85 except Exception as e:
86 print(f"WebSocket error: {e}")
87
88
89 def start_recording():
90 """Start the transcription service."""
91 global transcription_text
92 transcription_text = ""
93 thread = threading.Thread(target=start_websocket, daemon=True)
94 thread.start()
95 return gr.update(interactive=False), gr.update(interactive=True), ""
96
97
98 def stop_recording():
99 """Stop the transcription service."""
100 global is_running
101 is_running = False
102 return gr.update(interactive=True), gr.update(interactive=False), transcription_text
103
104
105 def process_audio(audio):
106 """Process incoming audio and queue for streaming."""
107 global transcription_text
108
109 if audio is None or not is_running:
110 return transcription_text
111
112 sample_rate, audio_data = audio
113
114 # Convert to mono if stereo
115 if len(audio_data.shape) > 1:
116 audio_data = audio_data.mean(axis=1)
117
118 # Normalize to float
119 if audio_data.dtype == np.int16:
120 audio_float = audio_data.astype(np.float32) / 32767.0
121 else:
122 audio_float = audio_data.astype(np.float32)
123
124 # Resample to 16kHz if needed
125 if sample_rate != SAMPLE_RATE:
126 num_samples = int(len(audio_float) * SAMPLE_RATE / sample_rate)
127 audio_float = np.interp(
128 np.linspace(0, len(audio_float) - 1, num_samples),
129 np.arange(len(audio_float)),
130 audio_float,
131 )
132
133 # Convert to PCM16 and base64 encode
134 pcm16 = (audio_float * 32767).astype(np.int16)
135 b64_chunk = base64.b64encode(pcm16.tobytes()).decode("utf-8")
136 audio_queue.put(b64_chunk)
137
138 return transcription_text
139
140
141 # Gradio interface
142 with gr.Blocks(title="Real-time Speech Transcription") as demo:
143 gr.Markdown("# Real-time Speech Transcription")
144 gr.Markdown("Click **Start** and speak into your microphone.")
145
146 with gr.Row():
147 start_btn = gr.Button("Start", variant="primary")
148 stop_btn = gr.Button("Stop", variant="stop", interactive=False)
149
150 audio_input = gr.Audio(sources=["microphone"], streaming=True, type="numpy")
151 transcription_output = gr.Textbox(label="Transcription", lines=5)
152
153 start_btn.click(
154 start_recording, outputs=[start_btn, stop_btn, transcription_output]
155 )
156 stop_btn.click(stop_recording, outputs=[start_btn, stop_btn, transcription_output])
157 audio_input.stream(
158 process_audio, inputs=[audio_input], outputs=[transcription_output]
159 )
160
161 if __name__ == "__main__":
162 parser = argparse.ArgumentParser(
163 description="Realtime WebSocket Transcription with Gradio"
164 )
165 parser.add_argument(
166 "--model",
167 type=str,
168 default="mistralai/Voxtral-Mini-4B-Realtime-2602",
169 help="Model that is served and should be pinged.",
170 )
171 parser.add_argument(
172 "--host", type=str, default="localhost", help="vLLM server host"
173 )
174 parser.add_argument("--port", type=int, default=8000, help="vLLM server port")
175 parser.add_argument(
176 "--share", action="store_true", help="Create public Gradio link"
177 )
178 args = parser.parse_args()
179
180 ws_url = f"ws://{args.host}:{args.port}/v1/realtime"
181 model = args.model
182 demo.launch(share=args.share)