Update app.py

app.py

@@ -14,37 +14,43 @@ processor = preloaded["processor"]
 
 st.title("Audio Inversion with HuggingFace & Streamlit")
 
+# If this is the first run, create a new session state attribute for uploaded file
+if 'uploaded_file' not in st.session_state:
+    st.session_state.uploaded_file = None
+
+# Get the uploaded file
 uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "flac"])
 
-if uploaded_file:
+# Update the session state only if a new file is uploaded
+if uploaded_file is not None:
+    st.session_state.uploaded_file = uploaded_file.getvalue()  # store content, not the file object
+
+if st.session_state.uploaded_file:
     # Play the uploaded audio
-    st.audio(uploaded_file, format="audio/wav")
+    audio_byte_content = st.session_state.uploaded_file
+    st.audio(audio_byte_content, format="audio/wav")
     
     # Read the audio file
-    audio, sr = sf.read(io.BytesIO(uploaded_file.getvalue()))
+    audio, sr = sf.read(io.BytesIO(audio_byte_content))
 
     # Convert audio to tensor
     audio_tensor = torch.tensor(audio).float()
 
-    # Use Streamlit's session state to prevent re-inversion
-    if "inverted_audio" not in st.session_state:
-        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()
+    @st.cache(allow_output_mutation=True, suppress_st_warning=True)
+    def cache_inverted_audio(audio_tensor):
+        return invert_audio(model, processor, audio_tensor, sr)
 
-            # Store inverted audio in session state
-            st.session_state.inverted_audio = inverted_audio_np
+    # Use cached result
+    inverted_audio_tensor = cache_inverted_audio(audio_tensor)
+    inverted_audio_np = inverted_audio_tensor.numpy()
 
-    # Play inverted audio from session state
+    # Play inverted audio
     with io.BytesIO() as out_io:
-        sf.write(out_io, st.session_state.inverted_audio, sr, format="wav")
+        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"):
         with io.BytesIO() as out_io:
-            sf.write(out_io, st.session_state.inverted_audio, sr, format="wav")
+            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")