Update app.py

app.py

@@ -94,46 +94,58 @@ from chromadb.config import Settings
 def extract_audio(uploaded_file):
     """Extracts audio from video or handles audio file directly."""
     audio_path = "temp_audio.wav"
-    
-    # Write uploaded file to a temporary file
+
+    # Save uploaded file temporarily
     temp_file = f"temp_{uploaded_file.name}"
     with open(temp_file, "wb") as f:
         f.write(uploaded_file.getvalue())
 
-    # Extract audio from video files or use directly for audio files
+    # Extract audio for video files, keep as-is for audio
     if uploaded_file.name.endswith(('.mp4', '.mkv')):
         ffmpeg.input(temp_file).output(audio_path).run(overwrite_output=True)
     else:
         audio_path = temp_file
-    
-    return audio_path
+
+    return audio_path, temp_file
 
 def transcribe_audio(audio_path):
-    """Transcribes audio to text using Whisper with model download handling."""
+    """Transcribes audio to text using Whisper."""
     try:
         model = whisper.load_model("base")
     except Exception:
         st.warning("Downloading Whisper model. This may take a while...")
         model = whisper.load_model("base")
+
     result = model.transcribe(audio_path)
     return result['text']
 
 def load_embeddings():
     """Loads subtitle embeddings from pkl file."""
-    with open('subtitle_embeddings.pkl', 'rb') as f:
-        embeddings = pickle.load(f)
-    return embeddings
+    if os.path.exists('subtitle_embeddings.pkl'):
+        with open('subtitle_embeddings.pkl', 'rb') as f:
+            embeddings = pickle.load(f)
+        return embeddings
+    else:
+        st.error("No embeddings file found.")
+        return pd.DataFrame()
 
 def save_to_chroma(embeddings):
     """Stores embeddings in Chroma DB."""
     client = Client(Settings())
-    collection = client.create_collection(name="subtitles")
+    
+    # Check if collection exists, else create
+    try:
+        collection = client.get_collection("subtitles")
+    except:
+        collection = client.create_collection(name="subtitles")
+    
     for idx, row in embeddings.iterrows():
         collection.add(
             documents=[row['subtitle']],
             ids=[str(idx)],
             embeddings=[row['embedding']]
         )
+    
     return collection
 
 def search_subtitles(query, collection):
@@ -149,11 +161,11 @@ def main():
         uploaded_file = st.file_uploader("Upload Video/Audio", type=["mp4", "mkv", "mp3", "wav"])
         query = st.text_input("Search Subtitles")
         download_btn = st.button("Download Subtitles")
-        
+
     if uploaded_file:
         with st.spinner("Extracting audio..."):
-            audio_path = extract_audio(uploaded_file.name)
-        
+            audio_path, temp_file = extract_audio(uploaded_file)
+
         with st.spinner("Generating subtitles..."):
             subtitles = transcribe_audio(audio_path)
             st.success("Subtitles Generated!")
@@ -163,23 +175,34 @@ def main():
             st.video(uploaded_file)
         else:
             st.audio(uploaded_file)
-        
+
         st.text_area("Generated Subtitles", subtitles, height=300)
-        
+
         # Load and search embeddings
         embeddings = load_embeddings()
-        collection = save_to_chroma(embeddings)
-        
-        if query:
-            results = search_subtitles(query, collection)
-            st.write("### Matching Subtitles:")
-            for sub in results:
-                st.write(f"- {sub}")
-        
+        if not embeddings.empty:
+            collection = save_to_chroma(embeddings)
+
+            if query:
+                results = search_subtitles(query, collection)
+                st.write("### Matching Subtitles:")
+                for sub in results:
+                    st.write(f"- {sub}")
+
+        # Subtitle download option
         if download_btn:
-            with open("generated_subtitles.srt", "w") as f:
-                f.write(subtitles)
-            st.download_button("Download SRT", "generated_subtitles.srt")
+            srt_content = f"1\n00:00:00,000 --> 00:00:10,000\n{subtitles}\n"
+            
+            st.download_button(
+                label="Download SRT",
+                data=srt_content.encode('utf-8'),
+                file_name="generated_subtitles.srt",
+                mime="text/plain"
+            )
+
+        # Cleanup temporary files
+        os.remove(audio_path)
+        os.remove(temp_file)
 
 if __name__ == '__main__':
     main()