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 math import uuid import re import time from urllib.parse import quote from datetime import timedelta # --- CRITICAL ENVIRONMENT FIXES --- if os.environ.get("OMP_NUM_THREADS", "").endswith("m"): os.environ["OMP_NUM_THREADS"] = "1" try: if "ffmpeg" in torchaudio.list_audio_backends(): torchaudio.set_audio_backend("ffmpeg") except Exception: pass if not hasattr(torch.load, "_is_patched"): _original_torch_load = torch.load def patched_torch_load(*args, **kwargs): kwargs['weights_only'] = False return _original_torch_load(*args, **kwargs) patched_torch_load._is_patched = True torch.load = patched_torch_load try: safe_list = [typing.Any, torch.nn.modules.container.ModuleList, np.dtype] 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) 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 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}") if not hasattr(np, 'NaN'): np.NaN = np.nan import whisperx # --- Configuration --- 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 # --- HELPER FUNCTIONS --- def clean_json_response(text): """Strips markdown code fences from Gemini output to prevent JSON errors.""" try: pattern = r"`{3}(?:json)?\s*(.*?)`{3}" match = re.search(pattern, text, re.DOTALL) if match: return match.group(1).strip() return text.strip() except Exception: return text def escape_xml(text): """Safely escapes special characters for XML to prevent DOM parser errors.""" if not text: return "" text = str(text) text = text.replace("&", "&") text = text.replace("<", "<") text = text.replace(">", ">") text = text.replace('"', """) text = text.replace("'", "'") return text def timecode_to_frames(tc, fps): """Converts HH:MM:SS:FF to absolute integer frames based on timebase.""" if not tc or not re.match(r"\d{2}:\d{2}:\d{2}[:\.]\d{2}", tc): return 0 parts = re.split(r'[:\.]', tc) h, m, s, f = map(int, parts) timebase = int(round(fps)) # 29.97 becomes 30, 23.98 becomes 24 return (h * 3600 * timebase) + (m * 60 * timebase) + (s * timebase) + f def frames_to_timecode(frames, fps): """Converts absolute integer frames back to HH:MM:SS:FF (NDF).""" timebase = int(round(fps)) hours = frames // (3600 * timebase) minutes = (frames // (60 * timebase)) % 60 secs = (frames // timebase) % 60 f = frames % timebase return f"{hours:02}:{minutes:02}:{secs:02}:{f:02}" def seconds_to_frames(seconds, fps): """Converts real seconds to integer frames.""" return int(round(seconds * fps)) def generate_cmx_edl(edl_title, segments, source_name, fps=25, start_offset_frames=0): """Constructs a CMX 3600 formatted EDL with source offset.""" edl_lines = [f"TITLE: {edl_title}", "FCM: NON-DROP FRAME\n"] rec_start_frames = 0 reel_id = source_name.replace(" ", "_")[:8] # EDL reel IDs are traditionally short max 8 chars for i, seg in enumerate(segments, 1): seg_type = seg.get('type', 'clip') is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic') is_graphic = seg_type == 'graphic' if is_vo or is_graphic: text_content = seg.get('text', 'Placeholder') default_dur = max(1.0, len(text_content.split()) / 2.5) if is_vo else 4.0 duration_sec = float(seg.get('duration', default_dur)) duration_frames = seconds_to_frames(duration_sec, fps) src_in_frames = 0 src_out_frames = duration_frames current_source = "GEN_VO" if is_vo else "GEN_GFX" note_text = f"{'VO SCRIPT' if is_vo else 'GRAPHIC'}: {text_content}" else: src_start_val = float(seg.get('src_start', seg.get('start', 0.0))) src_end_val = float(seg.get('src_end', seg.get('end', 0.0))) src_in_frames = seconds_to_frames(src_start_val, fps) + start_offset_frames src_out_frames = seconds_to_frames(src_end_val, fps) + start_offset_frames duration_frames = src_out_frames - src_in_frames current_source = reel_id note_text = seg.get('note', 'Clip') src_in = frames_to_timecode(src_in_frames, fps) src_out = frames_to_timecode(src_out_frames, fps) rec_in = frames_to_timecode(rec_start_frames, fps) rec_out = frames_to_timecode(rec_start_frames + duration_frames, fps) edl_lines.append(f"{i:03} {current_source:8} V C {src_in} {src_out} {rec_in} {rec_out}") if not is_vo and not is_graphic: edl_lines.append(f"* FROM CLIP NAME: {source_name}") edl_lines.append(f"* {note_text}\n") gap_frames = seconds_to_frames(float(seg.get('gap', 0.0)), fps) rec_start_frames += duration_frames + gap_frames return "\n".join(edl_lines) def generate_xml(sequence_name, segments, source_name, fps=25, start_offset_frames=0): """Constructs a Robust XML with source timecode offset for Premiere AND Resolve.""" master_id = "masterfile-1" timebase = int(round(fps)) is_ntsc = "TRUE" if fps % 1 != 0 else "FALSE" start_tc_string = frames_to_timecode(start_offset_frames, fps) encoded_name = quote(source_name) path_url = f"file://localhost/{encoded_name}" lines = [] lines.append('') lines.append('') lines.append('') lines.append('') lines.append(f'\t{escape_xml(sequence_name)}') lines.append('\t') lines.append(f'\t\t{timebase}') lines.append(f'\t\t{is_ntsc}') lines.append('\t') lines.append('\t') # --- VIDEO TRACK --- lines.append('\t\t') # --- AUDIO TRACK --- lines.append('\t\t') lines.append('\t') lines.append('') lines.append('') return "\n".join(lines) def generate_transcript_txt(sequence_name, segments, fps=25, start_offset_frames=0, transcript_data=None): """Generates a human-readable text document of the edit decisions and exact transcription.""" lines = [f"Sequence Transcript: {sequence_name}", "=" * 50, ""] for i, seg in enumerate(segments, 1): seg_type = seg.get('type', 'clip') is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic') is_graphic = seg_type == 'graphic' if is_vo: vo_text = seg.get('text', 'VO') duration = float(seg.get('duration', max(1.0, len(vo_text.split()) / 2.5))) lines.append(f"Clip {i:02} | [GENERATED VOICE-OVER] | Est. Duration: {duration:.1f}s") lines.append(f"Script: \"{vo_text}\"") elif is_graphic: gfx_text = seg.get('text', 'Graphic Text') duration = float(seg.get('duration', 4.0)) lines.append(f"Clip {i:02} | [GRAPHIC CARD] | Est. Duration: {duration:.1f}s") lines.append(f"Text: \"{gfx_text}\"") else: src_start_val = float(seg.get('src_start', seg.get('start', 0.0))) src_end_val = float(seg.get('src_end', seg.get('end', 0.0))) src_in_frames = seconds_to_frames(src_start_val, fps) + start_offset_frames src_out_frames = seconds_to_frames(src_end_val, fps) + start_offset_frames src_in = frames_to_timecode(src_in_frames, fps) src_out = frames_to_timecode(src_out_frames, fps) note = seg.get('note', 'No note provided.') actual_text = "" if transcript_data: clip_words = [] for t_seg in transcript_data: if max(src_start_val, t_seg['start']) < min(src_end_val, t_seg['end']): clip_words.append(t_seg['text']) actual_text = " ".join(clip_words).strip() lines.append(f"Clip {i:02} | IN: {src_in} --> OUT: {src_out}") lines.append(f"AI Summary: {note}") if actual_text: lines.append(f"Exact Audio Text: \"{actual_text}\"") lines.append("") return "\n".join(lines) def call_gemini_for_edl(transcript_data, story_prompt, api_key, previous_edit=None): if not api_key: st.error("Gemini API Key is missing.") return None 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.\n\n" "NEW CAPABILITIES:\n" "1. You can generate Voice-Over (VO) to bridge gaps, introduce topics, or summarize.\n" "2. You can generate Graphic Cards (Title Cards) to display text on screen (e.g., location, date, or chapter title).\n" "CRITICAL: ONLY generate 'vo' or 'graphic' segments if the user explicitly requests voice over, narration, or graphic/title cards in their creative brief. Otherwise, ONLY use 'clip' segments.\n\n" "Output ONLY a valid JSON array of segments. Every segment MUST be a 'clip', a 'vo', or a 'graphic'.\n\n" "For 'clip' segments (extracting from the subject):\n" "{\"type\": \"clip\", \"src_start\": 12.5, \"src_end\": 25.0, \"note\": \"Subject talks about X\", \"gap\": 0.0}\n" "- IGNORE ALL INTERVIEWER COMMENTS.\n" "- REMOVE FLUFF: Delete 'um', 'ah', repeats, and irrelevant filler.\n" "- TIMESTAMP INTEGRITY: Use only the exact word-level start and end times from the data.\n\n" "For 'vo' segments (generating new voice-over):\n" "{\"type\": \"vo\", \"text\": \"The journey began years ago...\", \"duration\": 3.5, \"gap\": 0.0}\n" "- DURATION: Estimate duration accurately based on reading speed (approx 2.5 words per second).\n" "- Keep VO concise and in the style of a documentary narrator.\n\n" "For 'graphic' segments (generating text on screen):\n" "{\"type\": \"graphic\", \"text\": \"Chapter 1: The Beginning\", \"duration\": 4.0, \"gap\": 0.0}\n\n" "PACING: Group related clips. Add a 'gap' (in seconds) between distinct ideas." ) if previous_edit: prompt_text = ( f"Transcript Data:\n{json.dumps(transcript_data)}\n\n" f"Previous Edit You Generated:\n{json.dumps(previous_edit)}\n\n" f"Director's Note for Revision: {story_prompt}\n\n" "Please output a NEW JSON array of segments applying these requested changes to the previous edit." ) else: prompt_text = f"Creative Brief: {story_prompt}\n\nTranscript Data:\n{json.dumps(transcript_data)}" payload = { "contents": [ { "role": "user", "parts": [{"text": prompt_text}] } ], "systemInstruction": { "parts": [{"text": system_prompt}] }, "generationConfig": { "responseMimeType": "application/json" } } headers = { "Content-Type": "application/json" } # Prioritizes Pro models, falls back to Flash. All "lite" versions explicitly removed. models_to_try = [ "gemini-3.1-pro-preview", "gemini-3-flash-preview", "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.0-flash" ] endpoints_to_try = ["v1beta", "v1"] last_error = "" for model_name in models_to_try: for api_version in endpoints_to_try: url = f"https://generativelanguage.googleapis.com/{api_version}/models/{model_name}:generateContent?key={api_key}" try: res = requests.post(url, json=payload, headers=headers) # Success if res.status_code == 200: result_json = res.json() raw_text = result_json['candidates'][0]['content']['parts'][0]['text'] cleaned_text = clean_json_response(raw_text) return json.loads(cleaned_text), model_name, api_version # Handle High Demand (503) or Rate Limits (429) elif res.status_code in [503, 429]: last_error = f"Model {model_name} busy ({res.status_code}) on {api_version}." st.warning(f"Debug: {last_error} Skipping to next model...") break # Handle endpoints that don't exist (404) or API rejections (400) else: try: error_details = res.json().get('error', {}).get('message', res.text) except Exception: error_details = res.text last_error = f"Model {model_name} rejected request ({res.status_code}) on {api_version}: {error_details}" st.warning(f"Debug: {last_error} Trying fallback endpoint/model...") continue except json.JSONDecodeError: last_error = f"Model {model_name} returned invalid JSON on {api_version}." st.warning(f"Debug: {last_error} Trying fallback endpoint/model...") continue except requests.exceptions.RequestException as e: last_error = f"Network error connecting to {model_name} on {api_version}: {e}" st.warning(f"Debug: {last_error} Trying fallback endpoint/model...") continue except Exception as e: last_error = f"Error with {model_name} on {api_version}: {e}" st.warning(f"Debug: {last_error} Trying fallback endpoint/model...") continue st.error(f"All Gemini models failed or are currently experiencing high demand. Last error: {last_error}") return None, None, None # --- Streamlit UI --- st.set_page_config(page_title="Junior Editor", layout="wide") st.title("Junior Editor") st.markdown(""" **Instructions** * Upload your file here (video or audio). * Set your timeline FPS, Source Timecode, and transcription quality in the sidebar. * Junior Editor will transcribe and separate speakers. You can then instruct it to find engaging bits or construct a narrative. * **Note: If you want me to write voice over suggestions or title cards, please request it.** * It will create an EDL or XML to import back into your editing software (Resolve, Premiere). * **For Multicam Workflows:** Use the **XML** option in the sidebar and enter the **exact name** of your Multicam Sequence. * **Resolve Users:** Uncheck "Automatically import source clips into media pool" during import. * **Premiere Users:** Use the EDL option if XML conformance fails for Multicam sequences. """) st.divider() with st.sidebar: clip_settings_container = st.container() transcription_settings_container = st.container() with clip_settings_container: st.header("Clip Settings") source_start_tc = st.text_input("Source Start Timecode", value="00:00:00:00", help="Format: HH:MM:SS:FF") export_format = st.radio("Output Format", ["EDL", "XML (Multicam)"], index=0) input_label = "EDL Reel Name" input_help = "Leave empty to use the uploaded file name." if export_format == "XML (Multicam)": input_label = "Multicam Sequence Name" input_help = "EXACT name of your Multicam Clip in Resolve." custom_reel_name = st.text_input( input_label, placeholder="e.g. Interview_Day1_Multi", help=input_help ) fps_options = [23.98, 24, 25, 29.97, 30, 50, 59.94, 60] fps = st.selectbox("Timeline FPS", fps_options, index=2) with transcription_settings_container: st.header("Transcription Settings") 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) uploaded_file = st.file_uploader("Upload Video/Audio Clip", type=["mp4", "m4a", "wav", "mp3", "mov"]) if uploaded_file: # --- Auto-Reset Logic --- 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 if "edit_generated" in st.session_state: del st.session_state.edit_generated 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 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.") # Phase 2 status_text.markdown("**Phase 2/4: Transcribing (Whisper)... This is the longest step.**") compute_type = "float16" if device == "cuda" else "int8" model = whisperx.load_model("large-v2", device, compute_type=compute_type) audio = whisperx.load_audio("temp_audio.wav") result = model.transcribe(audio, batch_size=4, language=target_language) del model gc.collect() torch.cuda.empty_cache() progress_bar.progress(50) # Phase 3 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 status_text.markdown("**Phase 4/4: Identifying Speakers...**") diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device) diarize_kwargs = {"min_speakers": num_speakers, "max_speakers": num_speakers} if num_speakers > 0 else {} 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 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: {e}") if os.path.exists("temp_input"): os.remove("temp_input") st.stop() if "transcript" in st.session_state: st.divider() with st.expander("Transcript Preview", expanded=False): 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. If you want me to write voice over suggestions or title cards, please request it.") if st.button("Generate Edit"): if not ACTIVE_GEMINI_KEY: st.error("Gemini API Key required.") else: with st.spinner("Junior Editor is thinking..."): final_source_name = custom_reel_name.strip() if custom_reel_name.strip() else uploaded_file.name offset_frames = timecode_to_frames(source_start_tc, fps) edl_segments, used_model, used_endpoint = call_gemini_for_edl(st.session_state.transcript, brief, ACTIVE_GEMINI_KEY) if edl_segments: safe_filename = "".join([c for c in final_source_name if c.isalnum() or c in (' ', '_', '-')]).strip() safe_filename = safe_filename.replace(' ', '_') if not safe_filename: safe_filename = "junior_editor_cut" # Generate the selected edit format (EDL/XML) if export_format == "EDL": final_output = generate_cmx_edl(final_source_name, edl_segments, final_source_name, fps, offset_frames) ext = "edl" else: final_output = generate_xml(final_source_name, edl_segments, final_source_name, fps, offset_frames) ext = "xml" # Generate the companion human-readable transcript transcript_txt = generate_transcript_txt(final_source_name, edl_segments, fps, offset_frames, st.session_state.transcript) # Save to session state so buttons don't disappear on click st.session_state.edit_generated = True st.session_state.revision_count = 1 st.session_state.base_filename = safe_filename st.session_state.final_output = final_output st.session_state.ext = ext st.session_state.transcript_txt = transcript_txt st.session_state.safe_filename = safe_filename st.session_state.used_model = used_model st.session_state.used_endpoint = used_endpoint st.session_state.edl_segments = edl_segments # Render the download buttons outside the Generate block using session state if st.session_state.get("edit_generated"): st.subheader("Ready for Import") st.success(f"✅ Edit generated successfully using **{st.session_state.used_model}** (via {st.session_state.used_endpoint}).") # Display downloads in a neat row col1, col2 = st.columns(2) with col1: st.download_button(f"Download .{st.session_state.ext.upper()} Sequence", data=st.session_state.final_output, file_name=f"{st.session_state.safe_filename}.{st.session_state.ext}") with col2: st.download_button(f"Download Reference Transcript (.TXT)", data=st.session_state.transcript_txt, file_name=f"{st.session_state.safe_filename}-Transcript.txt") st.info("📝 **Note:** There seems to be a bug in Resolve that the first time you import the sequence into a bin, sometimes it won't relink. Just Delete the sequence and Import again and it should work.") st.divider() st.subheader("Revise Edit") revision_brief = st.text_area("Want changes? Tell Junior Editor:", placeholder="e.g. Make it twice as long, add a graphic card for location, or write a VO intro.") if st.button("Apply Revisions"): with st.spinner("Junior Editor is revising the edit..."): final_source_name = custom_reel_name.strip() if custom_reel_name.strip() else uploaded_file.name offset_frames = timecode_to_frames(source_start_tc, fps) new_edl_segments, used_model, used_endpoint = call_gemini_for_edl( st.session_state.transcript, revision_brief, ACTIVE_GEMINI_KEY, previous_edit=st.session_state.edl_segments ) if new_edl_segments: # Increment version count st.session_state.revision_count = st.session_state.get("revision_count", 1) + 1 rev_seq_name = f"{final_source_name} V{st.session_state.revision_count}" if export_format == "EDL": final_output = generate_cmx_edl(rev_seq_name, new_edl_segments, final_source_name, fps, offset_frames) ext = "edl" else: final_output = generate_xml(rev_seq_name, new_edl_segments, final_source_name, fps, offset_frames) ext = "xml" transcript_txt = generate_transcript_txt(rev_seq_name, new_edl_segments, fps, offset_frames, st.session_state.transcript) st.session_state.final_output = final_output st.session_state.ext = ext st.session_state.transcript_txt = transcript_txt st.session_state.used_model = used_model st.session_state.used_endpoint = used_endpoint st.session_state.edl_segments = new_edl_segments # Update safe filename for the download buttons to show _V2, _V3, etc. st.session_state.safe_filename = f"{st.session_state.base_filename}_V{st.session_state.revision_count}" st.rerun()