Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import whisper
 from pytube import YouTube
+import os
 
 def get_audio(url):
     yt = YouTube(url)
@@ -40,6 +41,15 @@ def save_srt(transcript):
         f.write(transcript)
     return True
 
+def download_srt(transcript):
+    with open("done.txt", "w") as f:
+        f.write(transcript)
+    with open("done.txt", "r") as f:
+        srt = format_to_srt(f.read())
+    with open("transcript.srt", "w") as f:
+        f.write(srt)
+    return st.download_button(label="Download Transcript (.srt)", data="transcript.srt")
+
 langs = ["None"] + sorted(list(whisper.tokenizer.LANGUAGES.values()))
 model_size = list(whisper._MODELS.keys())
 
@@ -56,15 +66,7 @@ if st.button("Transcribe"):
     st.text_area("Transcription of the video", transcript)
     if format == ".srt":
         save_srt(transcript)
-        if st.button("Download Transcript"):
-            file_ = open("transcript.srt", "rb")
-            contents = file_.read()
-            file_.close()
-            st.download_button(
-                label="Download Transcript",
-                data=contents,
-                file_name="done.srt",
-                mime="text/srt"
-            )
-
+        download_srt(transcript)
 
+if os.path.exists("done.txt"):
+    os.remove("done.txt")