Spaces:
Running
Running
| from typing import Literal | |
| import gradio as gr | |
| from gradio_webrtc import WebRTC, StreamHandler, AdditionalOutputs | |
| from numpy import ndarray | |
| import sphn | |
| import websockets.sync.client | |
| import numpy as np | |
| class MoshiHandler(StreamHandler): | |
| def __init__(self, | |
| url: str, | |
| expected_layout: Literal['mono', 'stereo'] = "mono", | |
| output_sample_rate: int = 24000, | |
| output_frame_size: int = 480) -> None: | |
| self.url = url | |
| proto, without_proto = self.url.split('://', 1) | |
| if proto in ['ws', 'http']: | |
| proto = "ws" | |
| elif proto in ['wss', 'https']: | |
| proto = "wss" | |
| self._generator = None | |
| self.output_chunk_size = 1920 | |
| self.ws = None | |
| self.ws_url = f"{proto}://{without_proto}/api/chat" | |
| self.stream_reader = sphn.OpusStreamReader(output_sample_rate) | |
| self.stream_writer = sphn.OpusStreamWriter(output_sample_rate) | |
| self.all_output_data = None | |
| super().__init__(expected_layout, output_sample_rate, output_frame_size, | |
| input_sample_rate=24000) | |
| def receive(self, frame: tuple[int, ndarray]) -> None: | |
| if not self.ws: | |
| self.ws = websockets.sync.client.connect(self.ws_url) | |
| _, array = frame | |
| array = array.squeeze().astype(np.float32) / 32768.0 | |
| self.stream_writer.append_pcm(array) | |
| bytes = b"\x01" + self.stream_writer.read_bytes() | |
| self.ws.send(bytes) | |
| def generator(self): | |
| for message in self.ws: | |
| if len(message) == 0: | |
| yield None | |
| kind = message[0] | |
| if kind == 1: | |
| payload = message[1:] | |
| self.stream_reader.append_bytes(payload) | |
| pcm = self.stream_reader.read_pcm() | |
| if self.all_output_data is None: | |
| self.all_output_data = pcm | |
| else: | |
| self.all_output_data = np.concatenate((self.all_output_data, pcm)) | |
| while self.all_output_data.shape[-1] >= self.output_chunk_size: | |
| yield (self.output_sample_rate, self.all_output_data[: self.output_chunk_size].reshape(1, -1)) | |
| self.all_output_data = np.array(self.all_output_data[self.output_chunk_size :]) | |
| elif kind == 2: | |
| payload = message[1:] | |
| yield AdditionalOutputs(payload.decode()) | |
| def emit(self) -> tuple[int, ndarray] | None: | |
| if not self.ws: | |
| return | |
| if not self._generator: | |
| self._generator = self.generator() | |
| try: | |
| return next(self._generator) | |
| except StopIteration: | |
| self.reset() | |
| def reset(self) -> None: | |
| self._generator = None | |
| self.all_output_data = None | |
| def copy(self) -> StreamHandler: | |
| return MoshiHandler(self.url, | |
| self.expected_layout, | |
| self.output_sample_rate, self.output_frame_size) | |
| def shutdown(self) -> None: | |
| if self.ws: | |
| self.ws.close() | |
| with gr.Blocks() as demo: | |
| gr.HTML( | |
| """ | |
| <div style='text-align: center'> | |
| <h1> | |
| Talk To Moshi (Powered by WebRTC ⚡️) | |
| </h1> | |
| <p> | |
| Each conversation is limited to 90 seconds. Once the time limit is up you can rejoin the conversation. | |
| </p> | |
| </div> | |
| """ | |
| ) | |
| response = gr.State(value="") | |
| chatbot = gr.Chatbot(type="messages", value=[]) | |
| webrtc = WebRTC(label="Conversation", modality="audio", mode="send-receive", rtc_configuration=None) | |
| webrtc.stream(MoshiHandler("https://freddyaboulton-moshi-server.hf.space"), | |
| inputs=[webrtc, chatbot], outputs=[webrtc], time_limit=90) | |
| def on_text(state, response): | |
| print("response", response) | |
| return state + response | |
| def add_text(chat_history, response): | |
| if len(chat_history) == 0: | |
| chat_history.append({"role": "assistant", "content": response}) | |
| else: | |
| chat_history[-1]["content"] = response | |
| return chat_history | |
| webrtc.on_additional_outputs(on_text, | |
| inputs=[response], outputs=response, | |
| queue=True, | |
| show_progress="hidden") | |
| response.change(add_text, inputs=[chatbot, response], outputs=chatbot, | |
| queue=True, show_progress="hidden") | |
| demo.launch() | |