| """A simple web interactive chat demo based on gradio.""" |
|
|
| import os |
| import gradio as gr |
| from gradio_webrtc import WebRTC, AdditionalOutputs, ReplyOnPause |
| import base64 |
| import numpy as np |
| import requests |
| import io |
| from pydub import AudioSegment |
|
|
|
|
| API_URL = os.getenv("API_URL", None) |
| client = None |
|
|
| if API_URL is None: |
| from inference import OmniInference |
| omni_client = OmniInference('./checkpoint', 'cuda:0') |
| omni_client.warm_up() |
|
|
|
|
| OUT_CHUNK = 4096 |
| OUT_RATE = 24000 |
| OUT_CHANNELS = 1 |
|
|
|
|
| |
| account_sid = os.environ.get("TWILIO_ACCOUNT_SID") |
| auth_token = os.environ.get("TWILIO_AUTH_TOKEN") |
|
|
| if account_sid and auth_token: |
| from twilio.rest import Client |
| client = Client(account_sid, auth_token) |
|
|
| token = client.tokens.create() |
|
|
| rtc_configuration = { |
| "iceServers": token.ice_servers, |
| "iceTransportPolicy": "relay", |
| } |
| else: |
| rtc_configuration = None |
|
|
| |
| def response(audio: tuple[int, np.ndarray], conversation: list[dict], img: str | None): |
|
|
| sampling_rate, audio_np = audio |
| audio_np = audio_np.squeeze() |
|
|
| audio_buffer = io.BytesIO() |
| segment = AudioSegment( |
| audio_np.tobytes(), |
| frame_rate=sampling_rate, |
| sample_width=audio_np.dtype.itemsize, |
| channels=1, |
| ) |
|
|
| segment.export(audio_buffer, format="wav") |
| conversation.append({"role": "user", "content": gr.Audio((sampling_rate, audio_np))}) |
| conversation.append({"role": "assistant", "content": ""}) |
|
|
| base64_encoded = str(base64.b64encode(audio_buffer.getvalue()), encoding="utf-8") |
| if API_URL is not None: |
| output_audio_bytes = b"" |
| files = {"audio": base64_encoded} |
| if img is not None: |
| files["image"] = str(base64.b64encode(open(img, "rb").read()), encoding="utf-8") |
| print("sending request to server") |
| resp_text = "" |
| with requests.post(API_URL, json=files, stream=True) as response: |
| try: |
| buffer = b'' |
| for chunk in response.iter_content(chunk_size=2048): |
| buffer += chunk |
| while b'\r\n--frame\r\n' in buffer: |
| frame, buffer = buffer.split(b'\r\n--frame\r\n', 1) |
| if b'Content-Type: audio/wav' in frame: |
| audio_data = frame.split(b'\r\n\r\n', 1)[1] |
| |
| output_audio_bytes += audio_data |
| audio_array = np.frombuffer(audio_data, dtype=np.int16).reshape(1, -1) |
| yield (OUT_RATE, audio_array, "mono") |
| elif b'Content-Type: text/plain' in frame: |
| text_data = frame.split(b'\r\n\r\n', 1)[1].decode() |
| resp_text += text_data |
| if len(text_data) > 0: |
| conversation[-1]["content"] = resp_text |
| yield AdditionalOutputs(conversation) |
| except Exception as e: |
| raise Exception(f"Error during audio streaming: {e}") from e |
| |
|
|
| def main(port=None): |
|
|
| with gr.Blocks() as demo: |
| gr.HTML( |
| """ |
| <h1 style='text-align: center'> |
| Mini-Omni-2 Chat (Powered by WebRTC ⚡️) |
| </h1> |
| """ |
| ) |
| with gr.Row(): |
| with gr.Column(): |
| with gr.Row(): |
| with gr.Column(): |
| audio = WebRTC( |
| label="Stream", |
| rtc_configuration=rtc_configuration, |
| mode="send-receive", |
| modality="audio", |
| ) |
| with gr.Column(): |
| img = gr.Image(label="Image", type="filepath") |
| with gr.Column(): |
| conversation = gr.Chatbot(label="Conversation", type="messages") |
| |
| audio.stream( |
| fn=ReplyOnPause( |
| response, output_sample_rate=OUT_RATE, output_frame_size=480 |
| ), |
| inputs=[audio, conversation, img], |
| outputs=[audio], |
| time_limit=90, |
| ) |
| audio.on_additional_outputs(lambda c: c, outputs=[conversation]) |
| if port is not None: |
| demo.queue().launch(share=False, server_name="0.0.0.0", server_port=port) |
| else: |
| demo.queue().launch() |
|
|
|
|
| if __name__ == "__main__": |
| import fire |
|
|
| fire.Fire(main) |
|
|
