Update app.py

app.py

@@ -9,79 +9,86 @@ import os
 device = "cpu"
 torch_dtype = torch.float32
 
-# Load the Whisper model and processor
-model_id = "openai/whisper-large-v3-turbo"
-model = AutoModelForSpeechSeq2Seq.from_pretrained(
-    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
-).to(device)
-processor = AutoProcessor.from_pretrained(model_id)
+# Initialize session state
+if 'transcription_text' not in st.session_state:
+    st.session_state.transcription_text = None
+if 'srt_content' not in st.session_state:
+    st.session_state.srt_content = None
 
-# Create the pipeline for transcription
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    torch_dtype=torch_dtype,
-    device=device,
-)
+@st.cache_resource
+def load_model():
+    model_id = "openai/whisper-large-v3-turbo"
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+    ).to(device)
+    processor = AutoProcessor.from_pretrained(model_id)
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+    return pipe
+
+def format_srt_time(seconds):
+    hours, remainder = divmod(seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    milliseconds = int((seconds % 1) * 1000)
+    seconds = int(seconds)
+    return f"{int(hours):02}:{int(minutes):02}:{seconds:02},{milliseconds:03}"
 
 st.title("Audio/Video Transcription App")
 
+# Load model
+pipe = load_model()
+
 # File upload
 uploaded_file = st.file_uploader("Upload an audio or video file", type=["mp3", "wav", "mp4", "m4a"])
 
 if uploaded_file is not None:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
-        # If it's a video, extract audio
-        if uploaded_file.name.endswith(("mp4", "m4a")):
-            audio = AudioSegment.from_file(uploaded_file)
-            audio.export(temp_audio.name, format="wav")
-        else:
-            audio = AudioSegment.from_file(uploaded_file)
-            audio.export(temp_audio.name, format="wav")
-
-        # Run the transcription
-        transcription_result = pipe(temp_audio.name, return_timestamps="word")
-
-        # Extract text and timestamps
-        transcription_text = transcription_result['text']
-        transcription_chunks = transcription_result['chunks']
-
-        # Display transcription
-        st.subheader("Transcription")
-        st.write(transcription_text)
-
-        # Generate SRT file
-        srt_content = ""
-        for i, chunk in enumerate(transcription_chunks, start=1):
-            start_time = chunk["timestamp"][0]
-            end_time = chunk["timestamp"][1]
-            text = chunk["text"]
-
-            # Format time for SRT (hours, minutes, seconds, milliseconds)
-            def format_srt_time(seconds):
-                hours, remainder = divmod(seconds, 3600)
-                minutes, seconds = divmod(remainder, 60)
-                milliseconds = int((seconds % 1) * 1000)
-                seconds = int(seconds)
-                return f"{int(hours):02}:{int(minutes):02}:{seconds:02},{milliseconds:03}"
-
-            srt_content += f"{i}\n"
-            srt_content += f"{format_srt_time(start_time)} --> {format_srt_time(end_time)}\n"
-            srt_content += f"{text}\n\n"
+    with st.spinner("Processing audio..."):
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+            # If it's a video, extract audio
+            if uploaded_file.name.endswith(("mp4", "m4a")):
+                audio = AudioSegment.from_file(uploaded_file)
+                audio.export(temp_audio.name, format="wav")
+            else:
+                audio = AudioSegment.from_file(uploaded_file)
+                audio.export(temp_audio.name, format="wav")
+            
+            # Run the transcription
+            transcription_result = pipe(temp_audio.name, return_timestamps="word")
+            
+            # Extract text and timestamps
+            st.session_state.transcription_text = transcription_result['text']
+            transcription_chunks = transcription_result['chunks']
+            
+            # Generate SRT content
+            srt_content = ""
+            for i, chunk in enumerate(transcription_chunks, start=1):
+                start_time = chunk["timestamp"][0]
+                end_time = chunk["timestamp"][1]
+                text = chunk["text"]
+                
+                srt_content += f"{i}\n"
+                srt_content += f"{format_srt_time(start_time)} --> {format_srt_time(end_time)}\n"
+                srt_content += f"{text}\n\n"
+            
+            st.session_state.srt_content = srt_content
 
-        # Save the SRT file
-        srt_path = tempfile.mktemp(suffix=".srt")
-        with open(srt_path, "w") as srt_file:
-            srt_file.write(srt_content)
+# Display transcription
+if st.session_state.transcription_text:
+    st.subheader("Transcription")
+    st.write(st.session_state.transcription_text)
 
-        # Provide download for SRT file
+    # Provide download for SRT file
+    if st.session_state.srt_content:
         st.subheader("Download SRT File")
-        with open(srt_path, "rb") as file:
-            st.download_button(
-                label="Download SRT",
-                data=file,
-                file_name="transcription.srt",
-                mime="text/plain"
-            )
+        st.download_button(
+            label="Download SRT",
+            data=st.session_state.srt_content,
+            file_name="transcription.srt",
+            mime="text/plain"
+        )