Update app.py

app.py

@@ -1,25 +1,20 @@
 import streamlit as st
 import subprocess
-import whisper
 import json
 import os
 import requests
 import torch
-import torchaudio
+import whisperx
+import gc
 from datetime import timedelta
-from pyannote.audio import Pipeline
-from huggingface_hub import login, hf_hub_download
-from pydub import AudioSegment
 
 # --- Configuration & Tokens ---
-# Hardcode tokens here if you want to avoid UI input
 HARDCODED_HF_TOKEN = "PASTE_YOUR_HF_TOKEN_HERE"
 HARDCODED_GEMINI_KEY = ""
 
 ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "")
 ENV_GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")
 
-# Determine active keys
 ACTIVE_HF_TOKEN = ENV_HF_TOKEN if ENV_HF_TOKEN else HARDCODED_HF_TOKEN
 ACTIVE_GEMINI_KEY = ENV_GEMINI_KEY if ENV_GEMINI_KEY else HARDCODED_GEMINI_KEY
 
@@ -87,13 +82,17 @@ def call_gemini_for_edl(transcript_data, story_prompt, api_key):
         return None
 
 # --- Streamlit UI ---
-st.set_page_config(page_title="DocAI Editor", layout="wide")
-st.title("Documentary AI: Pipeline (Stable 2.1)")
+st.set_page_config(page_title="DocAI Editor (WhisperX)", layout="wide")
+st.title("Documentary AI: Pipeline (WhisperX)")
 
 with st.sidebar:
     st.header("Project Settings")
     fps = st.number_input("Timeline FPS", value=25)
     
+    st.header("Model Settings")
+    model_size = st.selectbox("Whisper Model Size", ["large-v3", "large-v2", "medium", "base"], index=0)
+    num_speakers = st.number_input("Number of Speakers (Optional)", min_value=0, value=0, help="Set to 0 for auto-detect")
+    
     st.divider()
     st.info("API Keys are managed via Environment Secrets.")
     if not ACTIVE_GEMINI_KEY:
@@ -110,105 +109,83 @@ if uploaded_file:
             if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
                 st.error("Please provide a valid Hugging Face Token.")
             else:
-                with st.spinner("Processing... This may take a moment."):
+                status_container = st.empty()
+                with status_container.container():
+                    st.write("🔄 **Processing Started...**")
+                    
                     # Save local temp file
                     with open("temp_input", "wb") as f:
                         f.write(uploaded_file.getbuffer())
                     
-                    st.write("🎵 **Preprocessing Audio...**")
-                    try:
-                        # Use PyDub to convert to WAV (Mono, 16kHz)
-                        audio = AudioSegment.from_file("temp_input")
-                        audio = audio.set_channels(1)
-                        audio = audio.set_frame_rate(16000)
-                        audio.export("temp_audio.wav", format="wav")
-                    except Exception as e:
-                        st.error(f"Audio processing failed: {e}")
-                        st.stop()
+                    st.write("🎵 **Extracting Audio (WAV)...**")
+                    # WhisperX prefers 16k mono wav
+                    subprocess.run([
+                        "ffmpeg", "-i", "temp_input", 
+                        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
+                        "temp_audio.wav", "-y"
+                    ])
                     
-                    # 1. Diarization (Pyannote 2.1.1)
-                    st.write("🗣️ **Running Speaker Diarization...**")
-                    diarization = None
                     try:
-                        # Log in globally
-                        login(token=ACTIVE_HF_TOKEN)
+                        device = "cuda" if torch.cuda.is_available() else "cpu"
+                        st.write(f"🚀 **Loading WhisperX on {device}...**")
                         
-                        # Use manual config download with explicit 'token' argument
-                        # This fixes the hf_hub_download error
-                        config_path = hf_hub_download(
-                            repo_id="pyannote/speaker-diarization", 
-                            revision="2.1",
-                            filename="config.yaml", 
-                            token=ACTIVE_HF_TOKEN
-                        )
+                        # 1. Transcribe with WhisperX (Faster-Whisper)
+                        batch_size = 16 # Reduce if low VRAM
+                        # Use float16 for GPU, int8 for CPU
+                        compute_type = "float16" if device == "cuda" else "int8"
                         
-                        # Load pipeline from the manually downloaded config
-                        # We do NOT pass a token here because the config file is local
-                        pipeline = Pipeline.from_pretrained(config_path)
+                        model = whisperx.load_model(model_size, device, compute_type=compute_type)
                         
-                        if torch.cuda.is_available():
-                            st.write("🚀 Using GPU for Diarization")
-                            pipeline.to(torch.device("cuda"))
-                            
-                        # Run pipeline directly on file path
-                        diarization = pipeline("temp_audio.wav")
+                        st.write("📝 **Transcribing...**")
+                        audio = whisperx.load_audio("temp_audio.wav")
+                        result = model.transcribe(audio, batch_size=batch_size)
                         
-                    except Exception as e:
-                        st.error(f"Diarization Error: {e}")
-                        diarization = None
-                    
-                    # 2. Whisper Transcription
-                    st.write("📝 **Transcribing with Whisper...**")
-                    device = "cuda" if torch.cuda.is_available() else "cpu"
-                    model = whisper.load_model("medium", device=device) 
-                    result = model.transcribe("temp_audio.wav", word_timestamps=True)
-                    
-                    # 3. Alignment
-                    st.write("🔗 **Aligning Speakers...**")
-                    final_segments = []
-                    speaker_turns = []
-                    
-                    if diarization:
-                        try:
-                            # 2.1.1 returns a proper Annotation object directly
-                            for turn, _, speaker_id in diarization.itertracks(yield_label=True):
-                                speaker_turns.append({"start": turn.start, "end": turn.end, "speaker": speaker_id})
-                            
-                            if len(speaker_turns) > 0:
-                                st.write(f"✅ Found {len(speaker_turns)} speaker turns.")
-                            else:
-                                st.warning("⚠️ Pipeline ran but returned no tracks.")
-                        except AttributeError:
-                             st.error("Could not iterate tracks. Output object format mismatch.")
-
-                    for segment in result['segments']:
-                        mid_time = (segment['start'] + segment['end']) / 2
-                        speaker = "Unknown"
+                        # Cleanup VRAM
+                        model_a = None
+                        gc.collect()
+                        torch.cuda.empty_cache()
+                        
+                        # 2. Align (Improves timestamp accuracy)
+                        st.write("⏱️ **Aligning Audio...**")
+                        model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
+                        result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
+                        
+                        # Cleanup VRAM
+                        model_a = None
+                        gc.collect()
+                        torch.cuda.empty_cache()
                         
-                        if speaker_turns:
-                            for turn in speaker_turns:
-                                if turn["start"] <= mid_time <= turn["end"]:
-                                    speaker = turn["speaker"]
-                                    break
+                        # 3. Diarize
+                        st.write("🗣️ **Diarizing Speakers...**")
+                        diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
+                        
+                        # Optional: Enforce speaker count
+                        diarize_kwargs = {}
+                        if num_speakers > 0:
+                            diarize_kwargs = {"min_speakers": num_speakers, "max_speakers": num_speakers}
                             
-                            if speaker == "Unknown":
-                                best_dist = 1.0 
-                                for turn in speaker_turns:
-                                    dist = min(abs(turn["start"] - mid_time), abs(turn["end"] - mid_time))
-                                    if dist < best_dist:
-                                        best_dist = dist
-                                        speaker = turn["speaker"]
+                        diarize_segments = diarize_model(audio, **diarize_kwargs)
+                        
+                        # 4. Assign Speakers to Words
+                        st.write("🔗 **Merging Transcripts...**")
+                        final_result = whisperx.assign_word_speakers(diarize_segments, result)
+                        
+                        # Format for Gemini
+                        processed_segments = []
+                        for segment in final_result["segments"]:
+                            processed_segments.append({
+                                "speaker": segment.get("speaker", "Unknown"),
+                                "text": segment["text"],
+                                "start": segment["start"],
+                                "end": segment["end"]
+                            })
+                        
+                        st.session_state.transcript = processed_segments
+                        st.success(f"Complete! Found {len(processed_segments)} segments.")
 
-                        final_segments.append({
-                            "speaker": speaker,
-                            "text": segment['text'],
-                            "start": segment['start'],
-                            "end": segment['end'],
-                            "words": segment.get('words', [])
-                        })
-                    
-                    st.session_state.transcript = final_segments
-                    st.success("Complete!")
+                    except Exception as e:
+                        st.error(f"Processing Error: {e}")
+                        st.stop()
 
     if "transcript" in st.session_state:
         st.divider()