Hugging Face's logo

Spaces:
freddyaboulton
/
mini-omni2-webrtc
Running on T4

App Files Community
mini-omni2-webrtc
File size: 4,487 Bytes 
54011d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9a16ab9
54011d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c34843b
54011d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b6d7cbe
54011d4
 
 
 
39d826e
bf0cc6a
54011d4
 
 
 
 
 
 
 
 
 
 
 
 
 
9a16ab9
 
 
 
 
 
 
 
 
 
54011d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
import logging
import base64
import io
import os
from threading import Thread

import gradio as gr
import numpy as np
import requests
from gradio_webrtc import ReplyOnPause, WebRTC, AdditionalOutputs
from pydub import AudioSegment
from twilio.rest import Client

from server import serve

logging.basicConfig(level=logging.WARNING)
file_handler = logging.FileHandler("gradio_webrtc.log")
file_handler.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
file_handler.setFormatter(formatter)
logger = logging.getLogger("gradio_webrtc")
logger.setLevel(logging.DEBUG)
logger.addHandler(file_handler)


IP = "0.0.0.0"
PORT = 60808

thread = Thread(target=serve, daemon=True)
thread.start()


API_URL = "http://0.0.0.0:60808/chat"

# Only needed if deploying on cloud provider
account_sid = os.environ.get("TWILIO_ACCOUNT_SID")
auth_token = os.environ.get("TWILIO_AUTH_TOKEN")

if account_sid and auth_token:
    client = Client(account_sid, auth_token)

    token = client.tokens.create()

    rtc_configuration = {
        "iceServers": token.ice_servers,
        "iceTransportPolicy": "relay",
    }
else:
    rtc_configuration = None

OUT_CHANNELS = 1
OUT_RATE = 24000
OUT_SAMPLE_WIDTH = 2
OUT_CHUNK = 20 * 4096

    
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]
                            # audio_data = base64.b64decode(audio_data)
                            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
                conversation[-1]["content"] = resp_text
                yield AdditionalOutputs(conversation)
            except Exception as e:
               raise Exception(f"Error during audio streaming: {e}") from e


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])


demo.launch()