app.py

import os
import numpy as np
import torch
import typing
import torchaudio
import streamlit as st
import subprocess
import json
import requests
import gc
import pandas as pd
from datetime import timedelta

# --- CRITICAL ENVIRONMENT FIXES ---
# 1. Fix for Hugging Face millicore OMP_NUM_THREADS error
if os.environ.get("OMP_NUM_THREADS", "").endswith("m"):
    os.environ["OMP_NUM_THREADS"] = "1"

# 2. Force Torchaudio Backend
try:
    if "ffmpeg" in torchaudio.list_audio_backends():
        torchaudio.set_audio_backend("ffmpeg")
except Exception:
    pass

# 3. GLOBAL PYTORCH SECURITY BYPASS (One-Time Patch)
if not hasattr(torch.load, "_is_patched"):
    print("DEBUG: Applying Monkeypatch to torch.load")
    _original_torch_load = torch.load

    def patched_torch_load(*args, **kwargs):
        # FORCE Disable security check
        kwargs['weights_only'] = False 
        return _original_torch_load(*args, **kwargs)
    
    # Mark as patched to prevent recursion loop on Streamlit reruns
    patched_torch_load._is_patched = True
    torch.load = patched_torch_load
else:
    print("DEBUG: torch.load is already patched. Skipping.")

# 4. EXPLICIT SAFE GLOBALS WHITELIST
try:
    safe_list = [
        typing.Any,
        torch.nn.modules.container.ModuleList,
        np.dtype,
    ]
    
    # NumPy internals
    if hasattr(np, '_core') and hasattr(np._core, 'multiarray'):
        safe_list.append(np._core.multiarray.scalar)
    elif hasattr(np, 'core') and hasattr(np.core, 'multiarray'):
        safe_list.append(np.core.multiarray.scalar)

    # OmegaConf
    try:
        from omegaconf.listconfig import ListConfig
        from omegaconf.dictconfig import DictConfig
        from omegaconf.base import ContainerMetadata, Metadata, Node
        safe_list.extend([ListConfig, DictConfig, ContainerMetadata, Metadata, Node])
    except ImportError:
        pass

    # Pyannote internals
    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
        safe_list.extend([Specifications, Problem, Resolution, Model, SpeakerDiarization])
    except ImportError:
        pass

    torch.serialization.add_safe_globals(safe_list)
except Exception as e:
    print(f"Safe Globals Warning: {e}")

# Fix NumPy 2.0+ attribute removal
if not hasattr(np, 'NaN'):
    np.NaN = np.nan

import whisperx

# --- 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

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
    minutes = (total_seconds % 3600) // 60
    secs = total_seconds % 60
    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
    
    # Sanitize source name for the Reel ID column
    reel_id = source_name.replace(" ", "_")
    
    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}  {reel_id}       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")
        
        # --- NEW LOGIC: Support Gaps ---
        # Add the duration of this clip AND any requested gap to the record timeline
        gap = seg.get('gap', 0.0)
        rec_start += duration + gap
        
    return "\n".join(edl_lines)

def call_gemini_for_edl(transcript_data, story_prompt, api_key):
    if not api_key:
        st.error("Gemini API Key is missing.")
        return None
        
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-09-2025:generateContent?key={api_key}"
    
    # --- UPDATED SYSTEM PROMPT ---
    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', 'note', and optionally 'gap'. "
        "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.\n"
        "5. PACING: Group related clips together so they flow seamlessly. HOWEVER, between distinct ideas or sections, "
        "add a 'gap': 1.0 property (float, in seconds) to the segment preceding the break. "
        "This will leave blank space on the timeline to signal a topic change."
    )
    
    prompt_text = f"Creative Brief: {story_prompt}\n\nTranscript Data:\n{json.dumps(transcript_data)}"
    
    payload = {
        "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

# --- Streamlit UI ---
st.set_page_config(page_title="Junior Editor", layout="wide")
st.title("Junior Editor")

st.markdown("""
**Instructions**
* Upload your file here, either an video file or audio. 
* Set your timeline FPS, choose the quality of your transcription and ,if you know the language, set this to speed up the transcription phase. 
* Junior Editor will transcribe it and separate by speakers and then await your instruction. You can ask it to find the most engaging bits and put them together from “Speaker 1”, remove all of Speaker 2, or construct a narrative around whatever idea you choose. 
* It will create an EDL to import back into your editing software. If you want it to reference a Master Clip rather than the proxy, type this in the side bar at the start. 
""")

st.divider()

with st.sidebar:
    st.header("Project Settings")
    fps = st.number_input("Timeline FPS", value=25)
    
    st.info("💡 **Conform Helper**")
    custom_reel_name = st.text_input(
        "EDL Reel Name", 
        placeholder="Paste Raw File Name Here...", 
        help="Leave empty to use the uploaded file name. Use this to link proxies to original camera files."
    )

    st.header("Model Settings")
    model_size = st.selectbox("Whisper Model", ["large-v2", "medium", "base"], index=0)
    
    language_map = {
        "Auto-Detect": None,
        "English": "en",
        "Spanish": "es",
        "French": "fr",
        "German": "de",
        "Italian": "it",
        "Portuguese": "pt"
    }
    selected_lang_label = st.selectbox("Audio Language", list(language_map.keys()), index=1)
    target_language = language_map[selected_lang_label]
    
    num_speakers = st.number_input("Speakers (0=Auto)", min_value=0, value=0)
    
    st.divider()
    if ACTIVE_HF_TOKEN == "PASTE_YOUR_HF_TOKEN_HERE":
        st.warning("⚠️ HF_TOKEN not set in Secrets!")
    else:
        st.success("✅ HF_TOKEN Loaded")

uploaded_file = st.file_uploader("Upload Video/Audio Clip", type=["mp4", "m4a", "wav", "mp3", "mov"])

if uploaded_file:
    # --- Auto-Reset Logic for New Files ---
    if "last_processed_file" not in st.session_state or st.session_state.last_processed_file != uploaded_file.name:
        if "transcript" in st.session_state:
            del st.session_state.transcript
        st.session_state.last_processed_file = uploaded_file.name

    # --- Auto-Process Logic ---
    if "transcript" not in st.session_state:
        if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
            st.error("Please provide a valid Hugging Face Token in the Sidebar/Secrets.")
        else:
            progress_container = st.container()
            with progress_container:
                st.info("🤖 **Junior Editor is processing your file...**")
                status_text = st.empty()
                progress_bar = st.progress(0)
                
                try:
                    # Phase 1: Save File
                    status_text.markdown("**Phase 1/4: Extracting Audio...**")
                    with open("temp_input", "wb") as f:
                        f.write(uploaded_file.getbuffer())
                    
                    subprocess.run([
                        "ffmpeg", "-i", "temp_input", 
                        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
                        "temp_audio.wav", "-y"
                    ])
                    progress_bar.progress(25)
                    
                    device = "cuda" if torch.cuda.is_available() else "cpu"
                    if device == "cpu":
                        st.warning("⚠️ No GPU detected. This will be slow.")
                    
                    # Phase 2: Transcribe
                    status_text.markdown(f"**Phase 2/4: Transcribing (Whisper {model_size})... This is the longest step.**")
                    
                    compute_type = "float16" if device == "cuda" else "int8"
                    model = whisperx.load_model(model_size, device, compute_type=compute_type)
                    audio = whisperx.load_audio("temp_audio.wav")
                    
                    result = model.transcribe(audio, batch_size=16, language=target_language)
                    
                    del model
                    gc.collect()
                    torch.cuda.empty_cache()
                    progress_bar.progress(50)
                    
                    # Phase 3: Align
                    status_text.markdown("**Phase 3/4: Aligning Text...**")
                    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)
                    
                    del model_a
                    gc.collect()
                    torch.cuda.empty_cache()
                    progress_bar.progress(75)
                    
                    # Phase 4: Diarize
                    status_text.markdown("**Phase 4/4: Identifying Speakers...**")
                    diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
                    
                    diarize_kwargs = {}
                    if num_speakers > 0:
                        diarize_kwargs = {"min_speakers": num_speakers, "max_speakers": num_speakers}
                    
                    diarize_segments = diarize_model(audio, **diarize_kwargs)
                    
                    # Final Merge
                    status_text.markdown("**Finalizing...**")
                    final_result = whisperx.assign_word_speakers(diarize_segments, result)
                    
                    processed_segments = []
                    for segment in final_result["segments"]:
                        processed_segments.append({
                            "speaker": segment.get("speaker", "Unknown"),
                            "text": segment["text"].strip(),
                            "start": segment["start"],
                            "end": segment["end"]
                        })
                    
                    st.session_state.transcript = processed_segments
                    
                    # --- CLEANUP ON SUCCESS ---
                    if os.path.exists("temp_input"): os.remove("temp_input")
                    if os.path.exists("temp_audio.wav"): os.remove("temp_audio.wav")
                    
                    progress_bar.progress(100)
                    status_text.success(f"Done! Processed {len(processed_segments)} segments.")

                except Exception as e:
                    status_text.error(f"Error during processing: {e}")
                    if os.path.exists("temp_input"): os.remove("temp_input")
                    if os.path.exists("temp_audio.wav"): os.remove("temp_audio.wav")
                    st.stop()

    if "transcript" in st.session_state:
        st.divider()
        with st.expander("Transcript Preview", expanded=True):
            for seg in st.session_state.transcript: 
                st.markdown(f"**{seg['speaker']}:** {seg['text']}")
        
        st.subheader("Your Instruction")
        brief = st.text_area("What should the Junior Editor do?", placeholder="e.g. Find the most engaging bits and put them together from Speaker 1, or remove all of Speaker 2.")
        
        if st.button("Generate Edit"):
            if not ACTIVE_GEMINI_KEY:
                st.error("Gemini API Key required in Secrets.")
            else:
                with st.spinner("Junior Editor is thinking..."):
                    final_source_name = custom_reel_name.strip() if custom_reel_name.strip() else uploaded_file.name
                    
                    edl_segments = call_gemini_for_edl(st.session_state.transcript, brief, ACTIVE_GEMINI_KEY)
                    if edl_segments:
                        final_edl = generate_cmx_edl("Junior_Editor_Cut", edl_segments, final_source_name, fps)
                        st.subheader("Ready for Import")
                        st.code(final_edl, language="text")
                        st.download_button("Download .EDL", data=final_edl, file_name="junior_editor_cut.edl")