Add streaming

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer
 import librosa
+from threading import Thread
 
 def split_audio(audio_arrays, chunk_limit=480000):
     CHUNK_LIM = chunk_limit
@@ -15,6 +16,14 @@ def split_audio(audio_arrays, chunk_limit=480000):
     return audio_splits
 
 
+def user(audio, text, chat_history):
+
+    if audio is not None:
+        chat_history.append(gr.ChatMessage(role="user", content={"path": audio, "alt_text": "Audio"}))
+    chat_history.append({"role": "user", "content": text})
+    return "", chat_history
+
+
 # Placeholder for your actual LLM processing API call
 def process_audio(audio, text, chat_history):
     conversation = [
@@ -24,6 +33,7 @@ def process_audio(audio, text, chat_history):
             ],
         },
     ]
+    audio_path = audio
     audio = librosa.load(audio, sr=16000)[0]
 
     if audio is not None:
@@ -35,7 +45,7 @@ def process_audio(audio, text, chat_history):
                     "audio": "placeholder",
                 }
             )
-        chat_history.append({"role": "user", "content": gr.Audio(value=(16000, audio))})
+        # chat_history.append(gr.ChatMessage(role="user", content={"path": audio_path, "alt_text": "Audio"}))
     
     conversation[0]["content"].append(
         {
@@ -45,22 +55,26 @@ def process_audio(audio, text, chat_history):
     )
 
     chat_history.append({"role": "user", "content": text})
+    # Set up the streamer for token generation
+    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
     prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
     inputs = processor(text=prompt, audios=splitted_audio, sampling_rate=16000, return_tensors="pt", padding=True)
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
-    outputs = model.generate(**inputs, eos_token_id=151645, pad_token_id=151643, max_new_tokens=4096)
-
-    cont = outputs[:, inputs["input_ids"].shape[-1] :]
-
-    result = processor.batch_decode(cont, skip_special_tokens=True)[0]
-    chat_history.append(
-        {
-            "role": "assistant",
-            "content": result,
-        }
+    # Set up generation arguments including max tokens and streamer
+    generation_args = {
+        "max_new_tokens": 4096,
+        "streamer": streamer,
+        **inputs
+    }
+    # Start a separate thread for model generation to allow streaming output
+    thread = Thread(
+        target=model.generate,
+        kwargs=generation_args,
     )
-
-    return chat_history
+    thread.start()
+    for character in streamer:
+        chat_history[-1]['content'] += character
+        yield chat_history 
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🎙️ Aero-1-Audio")
@@ -90,6 +104,11 @@ with gr.Blocks() as demo:
         chatbot_clear = gr.ClearButton([text_input, audio_input, chatbot], value="Clear")
         chatbot_submit = gr.Button("Submit", variant="primary")
         chatbot_submit.click(
+            user,
+            inputs=[audio_input, text_input, chatbot],
+            outputs=[text_input, chatbot],
+            queue=False
+        ).then(
             process_audio,
             inputs=[audio_input, text_input, chatbot],
             outputs=[chatbot],