Orca!

app.py

@@ -1,3 +1,5 @@
+import io
+import os
 import time
 import traceback
 from dataclasses import dataclass, field
@@ -5,6 +7,7 @@ from dataclasses import dataclass, field
 import gradio as gr
 import librosa
 import numpy as np
+import pvorca
 import soundfile as sf
 import spaces
 import torch
@@ -12,7 +15,8 @@ import xxhash
 from datasets import Audio
 from transformers import AutoModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
-import io
+
+orca = pvorca.create(access_key=os.environ.get("ORCA_KEY"))
 
 if gr.NO_RELOAD:
     diva_model = AutoModel.from_pretrained(
@@ -48,10 +52,8 @@ def diva_audio(audio_input, do_sample=False, temperature=0.001, prev_outs=None):
 
 @dataclass
 class AppState:
-    stream: np.ndarray | None = None
-    sampling_rate: int = 0
-    stopped: bool = False
     conversation: list = field(default_factory=list)
+    stopped: bool = False
     model_outs: any = None
 
 
@@ -63,16 +65,16 @@ def process_audio(audio: tuple, state: AppState):
 def response(state: AppState, audio: tuple):
     if not audio:
         return AppState()
-    state.stream = audio[1]
-    state.sampling_rate = audio[0]
 
-    file_name = f"/tmp/{xxhash.xxh32(bytes(state.stream)).hexdigest()}.wav"
+    file_name = f"/tmp/{xxhash.xxh32(bytes(audio[1])).hexdigest()}.wav"
 
-    sf.write(file_name, state.stream, state.sampling_rate, format="wav")
+    sf.write(file_name, audio[1], audio[0], format="wav")
 
     state.conversation.append(
         {"role": "user", "content": {"path": file_name, "mime_type": "audio/wav"}}
     )
+    state.conversation.append({"role": "assistant", "content": ""})
+    yield state, state.conversation, None
     if spaces.config.Config.zero_gpu:
         if state.model_outs is not None:
             state.model_outs = tuple(
@@ -88,18 +90,23 @@ def response(state: AppState, audio: tuple):
         causal_outs = state.model_outs
     state.model_outs = None
     prev_outs = causal_outs
-    start = False
+    stream = orca.stream_open()
     for resp, outs in diva_audio(
-        (state.sampling_rate, state.stream),
+        (audio[0], audio[1]),
         prev_outs=(prev_outs if prev_outs is not None else None),
     ):
-        if not start:
-            state.conversation.append({"role": "assistant", "content": resp})
-            start = True
-        else:
-            state.conversation[-1]["content"] = resp
-        print(resp)
-        yield state, state.conversation
+        prev_resp = state.conversation[-1]["content"]
+        state.conversation[-1]["content"] = resp
+        pcm = stream.synthesize(resp[len(prev_resp) :])
+        audio_chunk = None
+        if pcm is not None:
+            mp3_io = io.BytesIO()
+            sf.write(
+                mp3_io, np.asarray(pcm).astype(np.int16), orca.sample_rate, format="mp3"
+            )
+            audio_chunk = mp3_io.getvalue()
+            mp3_io.close()
+        yield state, state.conversation, audio_chunk
 
     del outs.logits
     del outs.hidden_states
@@ -107,9 +114,21 @@ def response(state: AppState, audio: tuple):
         outs = tuple(
             tuple(vec.cpu().numpy() for vec in tup) for tup in outs.past_key_values
         )
+    audio_chunk = None
+    pcm = stream.flush()
+    if pcm is not None:
+        audio_chunk = np.asarray(pcm).tobytes()
+        mp3_io = io.BytesIO()
+        sf.write(
+            mp3_io, np.asarray(pcm).astype(np.int16), orca.sample_rate, format="mp3"
+        )
+        audio_chunk = mp3_io.getvalue()
+        mp3_io.close()
+    stream.close()
     yield (
         AppState(conversation=state.conversation, model_outs=outs),
         state.conversation,
+        audio_chunk,
     )
 
 
@@ -190,6 +209,8 @@ with gr.Blocks(theme=theme, js=js) as demo:
         )
     with gr.Row():
         chatbot = gr.Chatbot(label="Conversation", type="messages")
+    with gr.Row():
+        output_audio = gr.Audio(label="Output Audio", streaming=True, autoplay=True)
     state = gr.State(value=AppState())
     stream = input_audio.start_recording(
         process_audio,
@@ -197,15 +218,15 @@ with gr.Blocks(theme=theme, js=js) as demo:
         [input_audio, state],
     )
     respond = input_audio.stop_recording(
-        response, [state, input_audio], [state, chatbot]
+        response, [state, input_audio], [state, chatbot, output_audio]
     )
-    restart = respond.success(start_recording_user, [state], [input_audio]).then(
+    restart = output_audio.stop(start_recording_user, [state], [input_audio]).then(
         lambda state: state, state, state, js=js_reset
     )
 
     cancel = gr.Button("Restart Conversation", variant="stop")
     cancel.click(
-        lambda: (AppState(stopped=True), gr.Audio(recording=False)),
+        lambda: (AppState(), gr.Audio(recording=False)),
         None,
         [state, input_audio],
         cancels=[respond, restart],