SpongeBobFan2002/openaudio-s1-mini

1import io
2import re
3import wave
4
5import gradio as gr
6import numpy as np
7
8from .fish_e2e import FishE2EAgent, FishE2EEventType
9from .schema import ServeMessage, ServeTextPart, ServeVQPart
10
11
12def wav_chunk_header(sample_rate=44100, bit_depth=16, channels=1):
13    buffer = io.BytesIO()
14
15    with wave.open(buffer, "wb") as wav_file:
16        wav_file.setnchannels(channels)
17        wav_file.setsampwidth(bit_depth // 8)
18        wav_file.setframerate(sample_rate)
19
20    wav_header_bytes = buffer.getvalue()
21    buffer.close()
22    return wav_header_bytes
23
24
25class ChatState:
26    def __init__(self):
27        self.conversation = []
28        self.added_systext = False
29        self.added_sysaudio = False
30
31    def get_history(self):
32        results = []
33        for msg in self.conversation:
34            results.append({"role": msg.role, "content": self.repr_message(msg)})
35
36        # Process assistant messages to extract questions and update user messages
37        for i, msg in enumerate(results):
38            if msg["role"] == "assistant":
39                match = re.search(r"Question: (.*?)\n\nResponse:", msg["content"])
40                if match and i > 0 and results[i - 1]["role"] == "user":
41                    # Update previous user message with extracted question
42                    results[i - 1]["content"] += "\n" + match.group(1)
43                    # Remove the Question/Answer format from assistant message
44                    msg["content"] = msg["content"].split("\n\nResponse: ", 1)[1]
45        return results
46
47    def repr_message(self, msg: ServeMessage):
48        response = ""
49        for part in msg.parts:
50            if isinstance(part, ServeTextPart):
51                response += part.text
52            elif isinstance(part, ServeVQPart):
53                response += f"<audio {len(part.codes[0]) / 21:.2f}s>"
54        return response
55
56
57def clear_fn():
58    return [], ChatState(), None, None, None
59
60
61async def process_audio_input(
62    sys_audio_input, sys_text_input, audio_input, state: ChatState, text_input: str
63):
64    if audio_input is None and not text_input:
65        raise gr.Error("No input provided")
66
67    agent = FishE2EAgent()  # Create new agent instance for each request
68
69    # Convert audio input to numpy array
70    if isinstance(audio_input, tuple):
71        sr, audio_data = audio_input
72    elif text_input:
73        sr = 44100
74        audio_data = None
75    else:
76        raise gr.Error("Invalid audio format")
77
78    if isinstance(sys_audio_input, tuple):
79        sr, sys_audio_data = sys_audio_input
80    else:
81        sr = 44100
82        sys_audio_data = None
83
84    def append_to_chat_ctx(
85        part: ServeTextPart | ServeVQPart, role: str = "assistant"
86    ) -> None:
87        if not state.conversation or state.conversation[-1].role != role:
88            state.conversation.append(ServeMessage(role=role, parts=[part]))
89        else:
90            state.conversation[-1].parts.append(part)
91
92    if state.added_systext is False and sys_text_input:
93        state.added_systext = True
94        append_to_chat_ctx(ServeTextPart(text=sys_text_input), role="system")
95    if text_input:
96        append_to_chat_ctx(ServeTextPart(text=text_input), role="user")
97        audio_data = None
98
99    result_audio = b""
100    async for event in agent.stream(
101        sys_audio_data,
102        audio_data,
103        sr,
104        1,
105        chat_ctx={
106            "messages": state.conversation,
107            "added_sysaudio": state.added_sysaudio,
108        },
109    ):
110        if event.type == FishE2EEventType.USER_CODES:
111            append_to_chat_ctx(ServeVQPart(codes=event.vq_codes), role="user")
112        elif event.type == FishE2EEventType.SPEECH_SEGMENT:
113            append_to_chat_ctx(ServeVQPart(codes=event.vq_codes))
114            yield state.get_history(), wav_chunk_header() + event.frame.data, None, None
115        elif event.type == FishE2EEventType.TEXT_SEGMENT:
116            append_to_chat_ctx(ServeTextPart(text=event.text))
117            yield state.get_history(), None, None, None
118
119    yield state.get_history(), None, None, None
120
121
122async def process_text_input(
123    sys_audio_input, sys_text_input, state: ChatState, text_input: str
124):
125    async for event in process_audio_input(
126        sys_audio_input, sys_text_input, None, state, text_input
127    ):
128        yield event
129
130
131def create_demo():
132    with gr.Blocks() as demo:
133        state = gr.State(ChatState())
134
135        with gr.Row():
136            # Left column (70%) for chatbot and notes
137            with gr.Column(scale=7):
138                chatbot = gr.Chatbot(
139                    [],
140                    elem_id="chatbot",
141                    bubble_full_width=False,
142                    height=600,
143                    type="messages",
144                )
145
146                # notes = gr.Markdown(
147                #     """
148                # # Fish Agent
149                # 1. 此Demo为Fish Audio自研端到端语言模型Fish Agent 3B版本.
150                # 2. 你可以在我们的官方仓库找到代码以及权重，但是相关内容全部基于 CC BY-NC-SA 4.0 许可证发布.
151                # 3. Demo为早期灰度测试版本，推理速度尚待优化.
152                # # 特色
153                # 1. 该模型自动集成ASR与TTS部分，不需要外挂其它模型，即真正的端到端，而非三段式(ASR+LLM+TTS).
154                # 2. 模型可以使用reference audio控制说话音色.
155                # 3. 可以生成具有较强情感与韵律的音频.
156                # """
157                # )
158                notes = gr.Markdown(
159                    """
160                    # Fish Agent
161                    1. This demo is Fish Audio's self-researh end-to-end language model, Fish Agent version 3B.
162                    2. You can find the code and weights in our official repo in [gitub](https://github.com/fishaudio/fish-speech) and [hugging face](https://huggingface.co/fishaudio/fish-agent-v0.1-3b), but the content is released under a CC BY-NC-SA 4.0 licence.
163                    3. The demo is an early alpha test version, the inference speed needs to be optimised.
164                    # Features
165                    1. The model automatically integrates ASR and TTS parts, no need to plug-in other models, i.e., true end-to-end, not three-stage (ASR+LLM+TTS).
166                    2. The model can use reference audio to control the speech timbre. 
167                    3. The model can generate speech with strong emotion.
168                """
169                )
170
171            # Right column (30%) for controls
172            with gr.Column(scale=3):
173                sys_audio_input = gr.Audio(
174                    sources=["upload"],
175                    type="numpy",
176                    label="Give a timbre for your assistant",
177                )
178                sys_text_input = gr.Textbox(
179                    label="What is your assistant's role?",
180                    value="You are a voice assistant created by Fish Audio, offering end-to-end voice interaction for a seamless user experience. You are required to first transcribe the user's speech, then answer it in the following format: 'Question: [USER_SPEECH]\n\nAnswer: [YOUR_RESPONSE]\n'. You are required to use the following voice in this conversation.",
181                    type="text",
182                )
183                audio_input = gr.Audio(
184                    sources=["microphone"], type="numpy", label="Speak your message"
185                )
186
187                text_input = gr.Textbox(label="Or type your message", type="text")
188
189                output_audio = gr.Audio(
190                    label="Assistant's Voice",
191                    streaming=True,
192                    autoplay=True,
193                    interactive=False,
194                )
195
196                send_button = gr.Button("Send", variant="primary")
197                clear_button = gr.Button("Clear")
198
199        # Event handlers
200        audio_input.stop_recording(
201            process_audio_input,
202            inputs=[sys_audio_input, sys_text_input, audio_input, state, text_input],
203            outputs=[chatbot, output_audio, audio_input, text_input],
204            show_progress=True,
205        )
206
207        send_button.click(
208            process_text_input,
209            inputs=[sys_audio_input, sys_text_input, state, text_input],
210            outputs=[chatbot, output_audio, audio_input, text_input],
211            show_progress=True,
212        )
213
214        text_input.submit(
215            process_text_input,
216            inputs=[sys_audio_input, sys_text_input, state, text_input],
217            outputs=[chatbot, output_audio, audio_input, text_input],
218            show_progress=True,
219        )
220
221        clear_button.click(
222            clear_fn,
223            inputs=[],
224            outputs=[chatbot, state, audio_input, output_audio, text_input],
225        )
226
227    return demo
228
229
230if __name__ == "__main__":
231    demo = create_demo()
232    demo.launch(server_name="127.0.0.1", server_port=7860, share=True)
233