Update app.py

app.py

@@ -8,15 +8,20 @@ import torch
 from datetime import timedelta
 from pyannote.audio import Pipeline
 from huggingface_hub import login
+from pydub import AudioSegment
 
-# ------------------ ENV SETUP ------------------
-os.environ["OMP_NUM_THREADS"] = "1"
+# --- Configuration & Tokens ---
+HARDCODED_HF_TOKEN = "PASTE_YOUR_HF_TOKEN_HERE"
+HARDCODED_GEMINI_KEY = ""
 
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
-GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")
+ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "")
+ENV_GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")
+
+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
 
-# ------------------ UTILS ------------------
 def format_timecode(seconds, fps=25):
+    """Converts seconds to HH:MM:SS:FF."""
     td = timedelta(seconds=seconds)
     total_seconds = int(td.total_seconds())
     hours = total_seconds // 3600
@@ -25,172 +30,186 @@ def format_timecode(seconds, fps=25):
     frames = int((seconds - total_seconds) * fps)
     return f"{hours:02}:{minutes:02}:{secs:02}:{frames:02}"
 
-def generate_cmx_edl(title, segments, source_name, fps=25):
-    edl = [f"TITLE: {title}", "FCM: NON-DROP FRAME\n"]
-    rec_time = 0.0
-
+def generate_cmx_edl(edl_title, segments, source_name, fps=25):
+    """Constructs a CMX 3600 formatted EDL."""
+    edl_lines = [f"TITLE: {edl_title}", "FCM: NON-DROP FRAME\n"]
+    rec_start = 0.0
     for i, seg in enumerate(segments, 1):
-        dur = seg["src_end"] - seg["src_start"]
-        edl.append(
-            f"{i:03}  AX       V     C        "
-            f"{format_timecode(seg['src_start'], fps)} "
-            f"{format_timecode(seg['src_end'], fps)} "
-            f"{format_timecode(rec_time, fps)} "
-            f"{format_timecode(rec_time + dur, fps)}"
-        )
-        edl.append(f"* FROM CLIP NAME: {source_name}")
-        edl.append(f"* {seg.get('note', '')}\n")
-        rec_time += dur
-
-    return "\n".join(edl)
-
-def call_gemini(transcript, brief, api_key):
+        src_in = format_timecode(seg['src_start'], fps)
+        src_out = format_timecode(seg['src_end'], fps)
+        duration = seg['src_end'] - seg['src_start']
+        rec_in = format_timecode(rec_start, fps)
+        rec_out = format_timecode(rec_start + duration, fps)
+        
+        edl_lines.append(f"{i:03}  AX       V     C        {src_in} {src_out} {rec_in} {rec_out}")
+        edl_lines.append(f"* FROM CLIP NAME: {source_name}")
+        edl_lines.append(f"* {seg.get('note', 'Clip')}\n")
+        rec_start += duration
+    return "\n".join(edl_lines)
+
+def call_gemini_for_edl(transcript_data, story_prompt, api_key):
+    """Sends diarized, word-level transcript to Gemini Senior Editor."""
     if not api_key:
+        st.error("Gemini API Key is missing. Set it in Space Secrets or app.py.")
         return None
-
-    url = (
-        "https://generativelanguage.googleapis.com/v1beta/"
-        "models/gemini-2.5-flash-preview-09-2025:generateContent"
-        f"?key={api_key}"
-    )
-
+        
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-09-2025:generateContent?key={api_key}"
+    
     system_prompt = (
-        "You are a documentary senior editor. "
-        "Using the diarized transcript JSON, produce a concise story cut. "
-        "Output ONLY a JSON array with src_start, src_end, note. "
-        "Ignore interviewer speech."
+        "You are an expert Documentary Senior Editor. Use the provided transcript JSON "
+        "(which includes Speaker IDs and word-level timestamps) to create a condensed story. "
+        "Output ONLY a valid JSON array of segments with 'src_start', 'src_end', and 'note'. "
+        "CRITICAL RULES:\n"
+        "1. IGNORE ALL INTERVIEWER COMMENTS: Do not include any speech or segments where the interviewer is speaking.\n"
+        "2. REMOVE FLUFF: Delete 'um', 'ah', repeats, and irrelevant filler.\n"
+        "3. NARRATIVE FLOW: Focus on the subject's high-energy responses and narrative hooks.\n"
+        "4. TIMESTAMP INTEGRITY: Use only the exact word-level start and end times from the data."
     )
-
+    
+    prompt_text = f"Creative Brief: {story_prompt}\n\nTranscript Data:\n{json.dumps(transcript_data)}"
+    
     payload = {
-        "contents": [{
-            "parts": [{"text": f"Brief:\n{brief}\n\nTranscript:\n{json.dumps(transcript)}"}]
-        }],
+        "contents": [{"parts": [{"text": prompt_text}]}],
         "systemInstruction": {"parts": [{"text": system_prompt}]},
         "generationConfig": {"responseMimeType": "application/json"}
     }
+    
+    try:
+        res = requests.post(url, json=payload)
+        res.raise_for_status()
+        result_json = res.json()
+        return json.loads(result_json['candidates'][0]['content']['parts'][0]['text'])
+    except Exception as e:
+        st.error(f"Senior Editor AI Error: {e}")
+        return None
 
-    r = requests.post(url, json=payload)
-    r.raise_for_status()
-    data = r.json()
-    return json.loads(data["candidates"][0]["content"]["parts"][0]["text"])
-
-# ------------------ UI ------------------
-st.set_page_config("DocAI Editor", layout="wide")
-st.title("🎬 Documentary AI Pipeline")
-
-if not HF_TOKEN:
-    st.error("HF_TOKEN missing in Space secrets")
-    st.stop()
-
-login(token=HF_TOKEN)
-
-uploaded = st.file_uploader(
-    "Upload audio or video",
-    type=["wav", "mp3", "m4a", "mp4", "mov"]
-)
-
-if uploaded:
-    if st.button("Step 1: Transcribe & Diarize"):
-        with st.spinner("Processing audio…"):
-
-            with open("input_media", "wb") as f:
-                f.write(uploaded.getbuffer())
-
-            subprocess.run(
-                [
-                    "ffmpeg", "-y",
-                    "-i", "input_media",
-                    "-vn",
-                    "-ac", "1",
-                    "-ar", "16000",
-                    "-acodec", "pcm_s16le",
-                    "audio.wav",
-                ],
-                check=True
-            )
-
-        # -------- DIARIZATION --------
-        st.write("🗣️ Running speaker diarization…")
-
-        pipeline = Pipeline.from_pretrained(
-            "pyannote/speaker-diarization-3.1",
-            use_auth_token=HF_TOKEN
-        )
-
-        diarization = pipeline("audio.wav")
-
-        speaker_turns = []
-        for turn, _, speaker in diarization.itertracks(yield_label=True):
-            speaker_turns.append({
-                "start": turn.start,
-                "end": turn.end,
-                "speaker": speaker
-            })
-
-        st.success(f"Detected {len(set(t['speaker'] for t in speaker_turns))} speakers")
-
-        # -------- TRANSCRIPTION --------
-        st.write("📝 Transcribing with Whisper…")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        whisper_model = whisper.load_model("medium", device=device)
-
-        result = whisper_model.transcribe(
-            "audio.wav",
-            word_timestamps=True
-        )
-
-        # -------- ALIGNMENT --------
-        st.write("🔗 Aligning speakers…")
-        final_segments = []
-
-        for seg in result["segments"]:
-            mid = (seg["start"] + seg["end"]) / 2
-            speaker = "Unknown"
-
-            for t in speaker_turns:
-                if t["start"] <= mid <= t["end"]:
-                    speaker = t["speaker"]
-                    break
-
-            final_segments.append({
-                "speaker": speaker,
-                "text": seg["text"],
-                "start": seg["start"],
-                "end": seg["end"],
-                "words": seg.get("words", [])
-            })
-
-        st.session_state.transcript = final_segments
-        st.success("Pipeline complete!")
-
-# ------------------ OUTPUT ------------------
-if "transcript" in st.session_state:
-    st.subheader("Transcript Preview")
-    for seg in st.session_state.transcript[:15]:
-        st.markdown(f"**{seg['speaker']}**: {seg['text']}")
-
-    brief = st.text_area("Creative Brief")
-
-    if st.button("Step 2: Create EDL"):
-        if not GEMINI_KEY:
-            st.error("Missing GEMINI_API_KEY")
-        else:
-            with st.spinner("Creating edit…"):
-                edl_segments = call_gemini(
-                    st.session_state.transcript,
-                    brief,
-                    GEMINI_KEY
-                )
-
-                edl = generate_cmx_edl(
-                    "AI_EDIT",
-                    edl_segments,
-                    uploaded.name
-                )
-
-                st.code(edl)
-                st.download_button(
-                    "Download EDL",
-                    edl,
-                    "edit.edl"
-                )
+# --- Streamlit UI ---
+st.set_page_config(page_title="DocAI Editor", layout="wide")
+st.title("Documentary AI: Pipeline (Stable v2.1)")
+
+with st.sidebar:
+    st.header("Project Settings")
+    fps = st.number_input("Timeline FPS", value=25)
+    
+    st.divider()
+    st.info("API Keys are managed via Environment Secrets.")
+    if not ACTIVE_GEMINI_KEY:
+        st.error("⚠️ Gemini API Key not found!")
+    if not ACTIVE_HF_TOKEN:
+        st.error("⚠️ HF Token not found!")
+
+uploaded_file = st.file_uploader("Upload Video/Audio Clip", type=["mp4", "m4a", "wav", "mp3", "mov"])
+
+if uploaded_file:
+    if "transcript" not in st.session_state:
+        if st.button("Step 1: Transcribe & Diarize"):
+            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."):
+                    # 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)
+                        # This is safer than raw ffmpeg for Pyannote 2.1
+                        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()
+                    
+                    # 1. Diarization (Pyannote 2.1.1)
+                    st.write("🗣️ **Running Speaker Diarization...**")
+                    diarization = None
+                    try:
+                        # Simple, legacy loading method
+                        pipeline = Pipeline.from_pretrained(
+                            "pyannote/speaker-diarization@2.1",
+                            use_auth_token=ACTIVE_HF_TOKEN
+                        )
+                        
+                        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")
+                        
+                    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:
+                        # 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.")
+
+                    for segment in result['segments']:
+                        mid_time = (segment['start'] + segment['end']) / 2
+                        speaker = "Unknown"
+                        
+                        if speaker_turns:
+                            for turn in speaker_turns:
+                                if turn["start"] <= mid_time <= turn["end"]:
+                                    speaker = turn["speaker"]
+                                    break
+                            
+                            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"]
+
+                        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!")
+
+    if "transcript" in st.session_state:
+        st.divider()
+        with st.expander("Transcript Preview (Diarized)"):
+            for seg in st.session_state.transcript[:20]: 
+                st.markdown(f"**{seg['speaker']}:** {seg['text']}")
+        
+        brief = st.text_area("Creative Brief", placeholder="e.g. Focus on the yeast story.")
+        
+        if st.button("Step 2: Create EDL"):
+            if not ACTIVE_GEMINI_KEY:
+                st.error("Gemini API Key required.")
+            else:
+                with st.spinner("Analyzing..."):
+                    edl_segments = call_gemini_for_edl(st.session_state.transcript, brief, ACTIVE_GEMINI_KEY)
+                    if edl_segments:
+                        final_edl = generate_cmx_edl("AI_Senior_Editor_Cut", edl_segments, uploaded_file.name, fps)
+                        st.subheader("EDL Preview")
+                        st.code(final_edl, language="text")
+                        st.download_button("Download EDL", data=final_edl, file_name="edit.edl")