Update app.py

app.py

@@ -1,34 +1,45 @@
-# functional app
+import streamlit as st
+import torch
+import julius
+import soundfile as sf
+import io
+
 from model import load_model, invert_audio
-import os
 
+# Load the model and processor
 preloaded = {}
 preloaded["model"], preloaded["processor"] = load_model()
-
-# Co
-for input_audio_path in data_dir.glob('*.wav'):
-    print(os.path.basename(input_audio_path))
-    output_audio_path = os.path.join(output_dir, "inverted-" + os.path.basename(input_audio_path))
-    output_file = invert_audio(
-        preloaded["model"], preloaded["processor"],
-        input_audio_path, output_audio_path)
-
-
-# HuggingFace UI
-import streamlit as st
-import torch
-import julius
+model = preloaded["model"]
+processor = preloaded["processor"]
 
 st.title("Audio Inversion with HuggingFace & Streamlit")
 
 uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "flac"])
 
 if uploaded_file:
+    # Play the uploaded audio
     st.audio(uploaded_file, format="audio/wav")
-    with st.spinner("Inverting audio..."):
-        output_path = "inverted_output.wav" # Temporary output path. Consider using temp files or dynamic naming in production.
-        invert_audio(model, processor, uploaded_file, output_path)
-        st.audio(output_path, format="audio/wav")
+    
+    # Read the audio file
+    audio, sr = sf.read(io.BytesIO(uploaded_file.getvalue()))
+
+    # Convert audio to tensor
+    audio_tensor = torch.tensor(audio).float()
 
+    with st.spinner("Inverting audio..."):
+        # Invert the audio using the modified function
+        inverted_audio_tensor = invert_audio(model, processor, audio_tensor, sr)
+        
+        # Convert tensor back to numpy
+        inverted_audio_np = inverted_audio_tensor.numpy()
+
+        # Convert numpy audio to bytes and play
+        with io.BytesIO() as out_io:
+            sf.write(out_io, inverted_audio_np, sr, format="wav")
+            st.audio(out_io.getvalue(), format="audio/wav")
+
+    # Offer a download button for the inverted audio
     if st.button("Download Inverted Audio"):
-        st.download_button("Download Inverted Audio", data=output_path, file_name="inverted_output.wav", mime="audio/wav")
+        with io.BytesIO() as out_io:
+            sf.write(out_io, inverted_audio_np, sr, format="wav")
+            st.download_button("Download Inverted Audio", data=out_io.getvalue(), file_name="inverted_output.wav", mime="audio/wav")