add state management

app.py

@@ -3,6 +3,7 @@ import requests
 import numpy as np
 import os
 import base64
+from streamlit import session_state as st_state
 
 # Try to get API_URL from environment variables, if not found set to a default value
 try:
@@ -23,6 +24,13 @@ headers = {
     "Content-Type": "application/json"
 }
 
+# Initialize session state variables
+if 'audio' not in st_state:
+    st_state.audio = None
+
+if 'augmented_audio' not in st_state:
+    st_state.augmented_audio = None
+
 # Streamlit app title
 st.title("Songlabai")
 
@@ -55,9 +63,9 @@ if st.button("Generate Audio"):
         payload = {"inputs": {"prompt": prompt, "duration": duration, "track": audio_base64}}
     st.text("Generating audio...")
     response = requests.post(API_URL, headers=headers, json=payload)
-    audio = np.array(response.json()[0]['generated_audio'], dtype=np.float32)
+    st_state.audio = np.array(response.json()[0]['generated_audio'], dtype=np.float32)
     sample_rate = response.json()[0]['sample_rate']
-    st.audio(audio, format="audio/wav", sample_rate=sample_rate, start_time=0)
+    st.audio(st_state.audio, format="audio/wav", sample_rate=sample_rate, start_time=0)
 
 # Post-processing options
 st.sidebar.title("Post-processing Options")
@@ -86,30 +94,32 @@ if apply_stereo or reverse or change_speed or pitch_shift or add_reverb:
     st.text("Applying post-processing...")
 
     # Apply selected post-processing
+    st_state.augmented_audio = st_state.audio.copy()
+
     if apply_stereo:
         # Create a stereo effect by duplicating the audio and panning left and right
-        audio_left = audio
-        audio_right = audio
-        audio = np.stack([audio_left, audio_right], axis=-1)
+        audio_left = st_state.augmented_audio
+        audio_right = st_state.augmented_audio
+        st_state.augmented_audio = np.stack([audio_left, audio_right], axis=-1)
 
     if reverse:
         # Reverse the audio array
-        audio = np.flip(audio)
+        st_state.augmented_audio = np.flip(st_state.augmented_audio)
 
     if change_speed:
         # Change the speed by resampling the audio
         new_sample_rate = int(sample_rate * speed_factor)
-        audio = scipy.signal.resample(audio, int(len(audio) * speed_factor))
+        st_state.augmented_audio = scipy.signal.resample(st_state.augmented_audio, int(len(st_state.augmented_audio) * speed_factor))
 
     if pitch_shift:
         # Pitch shift using the Fourier shift method
         pitch_shift_factor = 2 ** (pitch_semitones / 12)
-        audio = scipy.signal.resample(audio, int(len(audio) / pitch_shift_factor))
+        st_state.augmented_audio = scipy.signal.resample(st_state.augmented_audio, int(len(st_state.augmented_audio) / pitch_shift_factor))
 
     if add_reverb:
         # Apply reverb using the convolution method
         reverb_ir = scipy.signal.exponential(reverb_room_scale, reverb_damping, reverb_wet_only)
-        audio = scipy.signal.fftconvolve(audio, reverb_ir)
+        st_state.augmented_audio = scipy.signal.fftconvolve(st_state.augmented_audio, reverb_ir)
 
     # Play the processed audio
-    st.audio(audio, format="audio/wav", sample_rate=sample_rate, start_time=0)
+    st.audio(st_state.augmented_audio, format="audio/wav", sample_rate=sample_rate, start_time=0)