Update app.py

app.py

@@ -12,6 +12,7 @@ from huggingface_hub import hf_hub_download, list_repo_files
 from pydub import AudioSegment
 import io
 import tempfile
+import pydub
 
 #tts cpu model
 tts_model_str = "en_us_hifi_jets_cpu.addon"
@@ -64,6 +65,20 @@ def combine_audio_files(audio_files):
         os.remove(audio_file)  # Remove temporary files
     return combined
 
+def postprocess(prediction_value):
+    if isinstance(prediction_value, pydub.AudioSegment):
+        # Convert AudioSegment to numpy array
+        samples = np.array(prediction_value.get_array_of_samples())
+        # If stereo, take the mean of both channels
+        if prediction_value.channels == 2:
+            samples = np.mean(samples.reshape(-1, 2), axis=1)
+        # Return as tuple (sample_rate, samples)
+        return (prediction_value.frame_rate, samples.astype(np.float32))
+    else:
+        raise ValueError(f"Cannot process {type(prediction_value)} as Audio")
+
+    
+
 #guardrail model
 guard_llm = "llama-3.1-8b-instant"
 
@@ -123,9 +138,7 @@ async def greet(product,description):
         audio_file = text_to_speech(a_list[0])
         audio_files.append(audio_file)
         final_audio = combine_audio_files(audio_files)
-        output_file = "final_output.mp3"
-        final_audio.export(output_file, format="mp3")
-        yield final_audio
+        yield postprocess(final_audio)
     else:
         audio_files = []
         output = llm.create_chat_completion(
@@ -147,9 +160,7 @@ async def greet(product,description):
                 audio_file = text_to_speech(delta.get('content', ''))
                 audio_files.append(audio_file)
         final_audio = combine_audio_files(audio_files)
-        output_file = "final_output.mp3"
-        final_audio.export(output_file, format="mp3")
-        yield final_audio
+        yield postprocess(final_audio)
 
 demo = gr.Interface(fn=greet, inputs=["text","text"], outputs=gr.Audio(), concurrency_limit=10)
 demo.launch()