Update app.py

app.py

@@ -5,44 +5,18 @@ import json
 import os
 import requests
 import torch
-import numpy as np
-import librosa
 from datetime import timedelta
 from pyannote.audio import Pipeline
-from huggingface_hub import login, hf_hub_download
-
-# --- Safe Globals for PyTorch 2.6+ ---
-try:
-    from pyannote.audio.core.task import Specifications, Problem, Resolution
-    from pyannote.audio.core.model import Model
-    from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
-    
-    torch.serialization.add_safe_globals([
-        torch.torch_version.TorchVersion,
-        Specifications,
-        Problem,
-        Resolution,
-        Model,
-        SpeakerDiarization,
-        np.dtype,
-        torch.nn.modules.container.ModuleList,
-        np.core.multiarray.scalar
-    ])
-except Exception as e:
-    print(f"Safe Globals Warning: {e}")
-
-# --- Configuration & Tokens ---
-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", "")
-
-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
+from huggingface_hub import login
 
+# ------------------ ENV SETUP ------------------
+os.environ["OMP_NUM_THREADS"] = "1"
+
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+GEMINI_KEY = os.environ.get("GEMINI_API_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
@@ -51,212 +25,172 @@ 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(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
+def generate_cmx_edl(title, segments, source_name, fps=25):
+    edl = [f"TITLE: {title}", "FCM: NON-DROP FRAME\n"]
+    rec_time = 0.0
+
     for i, seg in enumerate(segments, 1):
-        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."""
+        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):
     if not api_key:
-        st.error("Gemini API Key is missing. Set it in Space Secrets or app.py.")
         return None
-        
-    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-09-2025:generateContent?key={api_key}"
-    
+
+    url = (
+        "https://generativelanguage.googleapis.com/v1beta/"
+        "models/gemini-2.5-flash-preview-09-2025:generateContent"
+        f"?key={api_key}"
+    )
+
     system_prompt = (
-        "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."
+        "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."
     )
-    
-    prompt_text = f"Creative Brief: {story_prompt}\n\nTranscript Data:\n{json.dumps(transcript_data)}"
-    
+
     payload = {
-        "contents": [{"parts": [{"text": prompt_text}]}],
+        "contents": [{
+            "parts": [{"text": f"Brief:\n{brief}\n\nTranscript:\n{json.dumps(transcript)}"}]
+        }],
         "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
 
-# --- Streamlit UI ---
-st.set_page_config(page_title="DocAI Editor", layout="wide")
-st.title("Documentary AI: Pipeline")
-
-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:
-    # --- Step 1: Technical Processing ---
-    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."):
-                    # Authenticate Globally
-                    login(token=ACTIVE_HF_TOKEN)
-                    
-                    with open("temp_input", "wb") as f:
-                        f.write(uploaded_file.getbuffer())
-                    
-                    # Convert to strict WAV using FFmpeg
-                    subprocess.run([
-                        "ffmpeg", "-i", "temp_input", 
-                        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
-                        "temp_audio.wav", "-y"
-                    ])
-                    
-                    # 1. Diarization with Librosa (Robust Audio Loading)
-                    st.write("🗣️ **Running Speaker Diarization...**")
-                    diarization = None
-                    try:
-                        # Load Config Manually
-                        config_path = hf_hub_download(
-                            repo_id="pyannote/speaker-diarization-3.1", 
-                            filename="config.yaml", 
-                            token=ACTIVE_HF_TOKEN
-                        )
-                        pipeline = Pipeline.from_pretrained(config_path)
-                        
-                        # Load Audio with Librosa (More robust than torchaudio in some containers)
-                        y, sr = librosa.load("temp_audio.wav", sr=16000)
-                        waveform = torch.tensor(y).unsqueeze(0) # Add channel dim
-                        
-                        # Move to GPU if available
-                        if torch.cuda.is_available():
-                            st.write("🚀 Using GPU for Diarization")
-                            pipeline.to(torch.device("cuda"))
-                            waveform = waveform.to(torch.device("cuda"))
-                            
-                        # Run Pipeline on Tensor directly
-                        diarization_output = pipeline({"waveform": waveform, "sample_rate": 16000})
-                        
-                        # Handle Output Wrapper
-                        if isinstance(diarization_output, tuple):
-                            diarization = diarization_output[0]
-                        else:
-                            diarization = diarization_output
-                            
-                        # Extract Annotation
-                        if not hasattr(diarization, "itertracks"):
-                             if hasattr(diarization_output, "annotation"):
-                                 diarization = diarization_output.annotation
-                             elif hasattr(diarization_output, "get"):
-                                 diarization = diarization_output.get("annotation", diarization_output)
-
-                    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:
-                            iterator = None
-                            if hasattr(diarization, 'itertracks'):
-                                iterator = diarization.itertracks(yield_label=True)
-                            
-                            if iterator:
-                                for turn, _, speaker_id in iterator:
-                                    speaker_turns.append({"start": turn.start, "end": turn.end, "speaker": speaker_id})
-                                st.write(f"✅ Found {len(speaker_turns)} speaker turns.")
-                            else:
-                                st.warning("⚠️ Pipeline ran but returned no iterable tracks.")
-                        except Exception as e:
-                            st.error(f"Error iterating tracks: {e}")
-
-                    for segment in result['segments']:
-                        mid_time = (segment['start'] + segment['end']) / 2
-                        speaker = "Unknown"
-                        
-                        if speaker_turns:
-                            # Match speaker
-                            for turn in speaker_turns:
-                                if turn["start"] <= mid_time <= turn["end"]:
-                                    speaker = turn["speaker"]
-                                    break
-                            
-                            # Fallback distance matching
-                            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")
+    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"
+                )