File size: 5,987 Bytes
405ac70 eab6f19 405ac70 e45ab03 405ac70 eab6f19 405ac70 |
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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import asyncio
import base64
import os
import time
from io import BytesIO
import gradio as gr
import numpy as np
from google import genai
from gradio_webrtc import (
AsyncAudioVideoStreamHandler,
WebRTC,
async_aggregate_bytes_to_16bit,
VideoEmitType,
AudioEmitType,
get_twilio_turn_credentials,
)
from PIL import Image
def encode_audio(data: np.ndarray) -> dict:
"""Encode Audio data to send to the server"""
return {"mime_type": "audio/pcm", "data": base64.b64encode(data.tobytes()).decode("UTF-8")}
def encode_image(data: np.ndarray) -> dict:
with BytesIO() as output_bytes:
pil_image = Image.fromarray(data)
pil_image.save(output_bytes, "JPEG")
bytes_data = output_bytes.getvalue()
base64_str = str(base64.b64encode(bytes_data), "utf-8")
return {"mime_type": "image/jpeg", "data": base64_str}
class GeminiHandler(AsyncAudioVideoStreamHandler):
def __init__(
self, expected_layout="mono", output_sample_rate=24000, output_frame_size=480
) -> None:
super().__init__(
expected_layout,
output_sample_rate,
output_frame_size,
input_sample_rate=16000,
)
self.audio_queue = asyncio.Queue()
self.video_queue = asyncio.Queue()
self.quit = asyncio.Event()
self.session = None
self.last_frame_time = 0
def copy(self) -> "GeminiHandler":
return GeminiHandler(
expected_layout=self.expected_layout,
output_sample_rate=self.output_sample_rate,
output_frame_size=self.output_frame_size,
)
async def video_receive(self, frame: np.ndarray):
if self.session:
# send image every 1 second
if time.time() - self.last_frame_time > 1:
self.last_frame_time = time.time()
await self.session.send(encode_image(frame))
if self.latest_args[2] is not None:
await self.session.send(encode_image(self.latest_args[2]))
self.video_queue.put_nowait(frame)
async def video_emit(self) -> VideoEmitType:
return await self.video_queue.get()
async def connect(self, api_key: str):
if self.session is None:
client = genai.Client(api_key=api_key, http_options={"api_version": "v1alpha"})
config = {"response_modalities": ["AUDIO"]}
async with client.aio.live.connect(
model="gemini-2.0-flash-exp", config=config
) as session:
self.session = session
asyncio.create_task(self.receive_audio())
await self.quit.wait()
async def generator(self):
while not self.quit.is_set():
turn = self.session.receive()
async for response in turn:
if data := response.data:
yield data
async def receive_audio(self):
async for audio_response in async_aggregate_bytes_to_16bit(
self.generator()
):
self.audio_queue.put_nowait(audio_response)
async def receive(self, frame: tuple[int, np.ndarray]) -> None:
_, array = frame
array = array.squeeze()
audio_message = encode_audio(array)
if self.session:
await self.session.send(audio_message)
async def emit(self) -> AudioEmitType:
if not self.args_set.is_set():
await self.wait_for_args()
if self.session is None:
asyncio.create_task(self.connect(self.latest_args[1]))
array = await self.audio_queue.get()
return (self.output_sample_rate, array)
def shutdown(self) -> None:
self.quit.set()
self.connection = None
self.args_set.clear()
self.quit.clear()
css = """
#video-source {max-width: 600px !important; max-height: 600 !important;}
"""
with gr.Blocks(css=css) as demo:
gr.HTML(
"""
<div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
<div style="background-color: var(--block-background-fill); border-radius: 8px">
<img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
</div>
<div>
<h1>Gen AI SDK Voice Chat</h1>
<p>Speak with Gemini using real-time audio + video streaming</p>
<p>Powered by <a href="https://gradio.app/">Gradio</a> and <a href=https://freddyaboulton.github.io/gradio-webrtc/">WebRTC</a>⚡️</p>
<p>Get an API Key <a href="https://support.google.com/googleapi/answer/6158862?hl=en">here</a></p>
</div>
</div>
"""
)
with gr.Row() as api_key_row:
api_key = gr.Textbox(label="API Key", type="password", placeholder="Enter your API Key", value=os.getenv("GOOGLE_API_KEY"))
with gr.Row(visible=False) as row:
with gr.Column():
webrtc = WebRTC(
label="Video Chat",
modality="audio-video",
mode="send-receive",
elem_id="video-source",
rtc_configuration=get_twilio_turn_credentials(),
icon="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png",
pulse_color="rgb(35, 157, 225)",
icon_button_color="rgb(35, 157, 225)",
)
with gr.Column():
image_input = gr.Image(label="Image", type="numpy", sources=["upload", "clipboard"])
webrtc.stream(
GeminiHandler(),
inputs=[webrtc, api_key, image_input],
outputs=[webrtc],
time_limit=90,
concurrency_limit=2,
)
api_key.submit(
lambda: (gr.update(visible=False), gr.update(visible=True)),
None,
[api_key_row, row],
)
if __name__ == "__main__":
demo.launch()
|