Save audio and script in streamlit session state

demo/app.py

@@ -50,30 +50,19 @@ def load_text_to_speech_model_and_tokenizer():
     return load_parler_tts_model_and_tokenizer("parler-tts/parler-tts-mini-v1", "cpu")
 
 
-gen_pod_button = "generate_podcast_clicked"
-save_audio_button = "save_podcast_clicked"
-save_text_button = "save_script_clicked"
+script = "script"
+audio = "audio"
+gen_button = "generate podcast button"
+if script not in st.session_state:
+    st.session_state[script] = ""
+if audio not in st.session_state:
+    st.session_state.audio = []
+if gen_button not in st.session_state:
+    st.session_state[gen_button] = False
 
-if gen_pod_button not in st.session_state:
-    st.session_state[gen_pod_button] = False
-if save_audio_button not in st.session_state:
-    st.session_state[save_audio_button] = False
-if save_text_button not in st.session_state:
-    st.session_state[save_text_button] = False
 
-
-def click_generate_podcast():
-    st.session_state[gen_pod_button] = True
-    st.session_state[save_audio_button] = False
-    st.session_state[save_text_button] = False
-
-
-def click_save_audio():
-    st.session_state[save_audio_button] = True
-
-
-def click_save_text():
-    st.session_state[save_text_button] = True
+def gen_button_clicked():
+    st.session_state[gen_button] = True
 
 
 st.title("Document To Podcast")
@@ -129,17 +118,17 @@ def click_save_text():
 
     system_prompt = st.text_area("Podcast generation prompt", value=PODCAST_PROMPT)
 
-    if st.button("Generate Podcast", on_click=click_generate_podcast):
+    if st.button("Generate Podcast", on_click=gen_button_clicked):
         with st.spinner("Generating Podcast..."):
             text = ""
-            complete_script = ""
-            complete_audio = []
             for chunk in text_to_text_stream(
                 clean_text, text_model, system_prompt=system_prompt.strip()
             ):
                 text += chunk
                 if text.endswith("\n") and "Speaker" in text:
-                    st.write(text)
+                    st.session_state.script += text
+                    st.write(st.session_state.script)
+
                     speaker_id = re.search(r"Speaker (\d+)", text).group(1)
                     with st.spinner("Generating Audio..."):
                         speech = text_to_speech(
@@ -149,24 +138,22 @@ def click_save_text():
                             SPEAKER_DESCRIPTIONS[speaker_id],
                         )
                     st.audio(speech, sample_rate=speech_model.config.sampling_rate)
-                    complete_audio.append(speech)
-                    complete_script += text
+                    st.session_state.audio.append(speech)
                     text = ""
-            if st.session_state[gen_pod_button]:
-                if st.button("Save Podcast to audio file", on_click=click_save_audio):
-                    complete_audio = np.concatenate(complete_audio)
-                    save_waveform_as_file(
-                        waveform=complete_audio,
-                        sampling_rate=speech_model.config.sampling_rate,
-                        filename="podcast.wav",
-                    )
-                    st.markdown("Podcast saved to disk!")
-
-                if st.button(
-                    "Save Podcast script to text file", on_click=click_save_text
-                ):
-                    with open("script.txt", "w") as f:
-                        complete_script += "}"
-                        f.write(complete_script)
-
-                    st.markdown("Script saved to disk!")
+
+    if st.session_state[gen_button]:
+        if st.button("Save Podcast to audio file"):
+            st.session_state.audio = np.concatenate(st.session_state.audio)
+            save_waveform_as_file(
+                waveform=st.session_state.audio,
+                sampling_rate=speech_model.config.sampling_rate,
+                filename="podcast.wav",
+            )
+            st.markdown("Podcast saved to disk!")
+
+        if st.button("Save Podcast script to text file"):
+            with open("script.txt", "w") as f:
+                st.session_state.script += "}"
+                f.write(st.session_state.script)
+
+            st.markdown("Script saved to disk!")