Spaces:
Build error
Build error
| import gradio as gr | |
| import subprocess | |
| import os | |
| import tempfile | |
| import shutil | |
| from pathlib import Path | |
| import json | |
| import datetime | |
| import csv | |
| from pydub import AudioSegment | |
| import numpy as np | |
| import torch | |
| import gc | |
| from dotenv import load_dotenv | |
| # Load environment variables | |
| load_dotenv() | |
| # Import NeMo for transcription (you'll need to install: pip install nemo_toolkit[asr]) | |
| try: | |
| from nemo.collections.asr.models import ASRModel | |
| NEMO_AVAILABLE = True | |
| except ImportError: | |
| NEMO_AVAILABLE = False | |
| print("Warning: NeMo not available. Auto-transcription will be disabled.") | |
| class AutomatedSubtitleBurner: | |
| def __init__(self): | |
| self.temp_dir = tempfile.mkdtemp() | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.model = None | |
| # Load transcription model if available | |
| if NEMO_AVAILABLE: | |
| try: | |
| MODEL_NAME = os.getenv('MODEL_NAME') | |
| if MODEL_NAME: | |
| self.model = ASRModel.from_pretrained(model_name=MODEL_NAME) | |
| self.model.eval() | |
| except Exception as e: | |
| self.model = None | |
| def extract_audio_from_video(self, video_path): | |
| """Extract audio from video file for transcription""" | |
| try: | |
| audio_path = os.path.join(self.temp_dir, "extracted_audio.wav") | |
| # Use FFmpeg to extract audio | |
| cmd = [ | |
| 'ffmpeg', '-y', '-i', video_path, | |
| '-vn', # No video | |
| '-acodec', 'pcm_s16le', # Audio codec | |
| '-ar', '16000', # Sample rate | |
| '-ac', '1', # Mono | |
| audio_path | |
| ] | |
| subprocess.run(cmd, capture_output=True, check=True) | |
| return audio_path | |
| except Exception as e: | |
| print(f"Error extracting audio: {e}") | |
| return None | |
| def format_srt_time(self, seconds: float) -> str: | |
| """Converts seconds to SRT time format HH:MM:SS,mmm""" | |
| sanitized_total_seconds = max(0.0, seconds) | |
| delta = datetime.timedelta(seconds=sanitized_total_seconds) | |
| total_int_seconds = int(delta.total_seconds()) | |
| hours = total_int_seconds // 3600 | |
| remainder_seconds_after_hours = total_int_seconds % 3600 | |
| minutes = remainder_seconds_after_hours // 60 | |
| seconds_part = remainder_seconds_after_hours % 60 | |
| milliseconds = delta.microseconds // 1000 | |
| return f"{hours:02d}:{minutes:02d}:{seconds_part:02d},{milliseconds:03d}" | |
| def generate_srt_content(self, word_timestamps: list) -> str: | |
| """Generates SRT formatted string from word timestamps""" | |
| srt_content = [] | |
| for i, ts in enumerate(word_timestamps): | |
| start_time = self.format_srt_time(ts['start']) | |
| end_time = self.format_srt_time(ts['end']) | |
| text = ts['word'] | |
| srt_content.append(str(i + 1)) | |
| srt_content.append(f"{start_time} --> {end_time}") | |
| srt_content.append(text) | |
| srt_content.append("") | |
| return "\n".join(srt_content) | |
| def transcribe_audio(self, audio_path, progress_callback=None): | |
| """Transcribe audio to get word-level timestamps""" | |
| if not self.model or not NEMO_AVAILABLE: | |
| return None, "Transcription model not available" | |
| try: | |
| if progress_callback: | |
| progress_callback(0.1, "Loading audio...") | |
| # Load and preprocess audio | |
| audio = AudioSegment.from_file(audio_path) | |
| duration_sec = audio.duration_seconds | |
| if progress_callback: | |
| progress_callback(0.2, "Preprocessing audio...") | |
| # Ensure audio is in correct format | |
| if audio.frame_rate != 16000: | |
| audio = audio.set_frame_rate(16000) | |
| if audio.channels != 1: | |
| audio = audio.set_channels(1) | |
| # Save preprocessed audio | |
| processed_path = os.path.join(self.temp_dir, "processed_audio.wav") | |
| audio.export(processed_path, format="wav") | |
| if progress_callback: | |
| progress_callback(0.3, "Starting transcription...") | |
| # Configure model for long audio if needed | |
| long_audio_settings_applied = False | |
| if duration_sec > 480: # 8 minutes | |
| try: | |
| print("Applying long audio settings for transcription...") | |
| self.model.change_attention_model("rel_pos_local_attn", [256, 256]) | |
| self.model.change_subsampling_conv_chunking_factor(1) | |
| long_audio_settings_applied = True | |
| except Exception as e: | |
| print(f"Warning: Could not apply long audio settings: {e}") | |
| # Move model to appropriate device and precision | |
| self.model.to(self.device) | |
| self.model.to(torch.bfloat16) | |
| if progress_callback: | |
| progress_callback(0.5, "Transcribing (this may take a while)...") | |
| # Transcribe with timestamps | |
| output = self.model.transcribe([processed_path], timestamps=True) | |
| if progress_callback: | |
| progress_callback(0.8, "Processing transcription results...") | |
| if not output or not output[0] or not hasattr(output[0], 'timestamp'): | |
| return None, "Transcription failed - no output generated" | |
| # Get word-level timestamps | |
| word_timestamps = output[0].timestamp.get('word', []) | |
| if not word_timestamps: | |
| return None, "No word-level timestamps generated" | |
| # Generate SRT content | |
| srt_content = self.generate_srt_content(word_timestamps) | |
| if progress_callback: | |
| progress_callback(1.0, "Transcription complete!") | |
| return srt_content, "Transcription successful!" | |
| except torch.cuda.OutOfMemoryError: | |
| return None, "CUDA out of memory. Please try a shorter video or use CPU." | |
| except Exception as e: | |
| return None, f"Transcription error: {str(e)}" | |
| finally: | |
| # Cleanup model settings and memory | |
| try: | |
| if long_audio_settings_applied and self.model: | |
| self.model.change_attention_model("rel_pos") | |
| self.model.change_subsampling_conv_chunking_factor(-1) | |
| if self.model and self.device == 'cuda': | |
| self.model.cpu() | |
| gc.collect() | |
| if self.device == 'cuda': | |
| torch.cuda.empty_cache() | |
| except Exception as e: | |
| print(f"Warning: Error during cleanup: {e}") | |
| def auto_generate_srt(self, video_file, progress=gr.Progress()): | |
| """Automatically generate SRT from video""" | |
| if not video_file: | |
| return "", "Please provide a video file" | |
| if not self.model or not NEMO_AVAILABLE: | |
| return "", "Transcription model not available. Please install NeMo toolkit." | |
| try: | |
| progress(0.05, desc="Extracting audio from video...") | |
| # Extract audio from video | |
| audio_path = self.extract_audio_from_video(video_file) | |
| if not audio_path: | |
| return "", "Failed to extract audio from video" | |
| progress(0.1, desc="Audio extracted, starting transcription...") | |
| # Transcribe audio | |
| def progress_callback(value, desc): | |
| progress(0.1 + (value * 0.8), desc=desc) | |
| srt_content, message = self.transcribe_audio(audio_path, progress_callback) | |
| progress(0.95, desc="Finalizing...") | |
| if srt_content: | |
| progress(1.0, desc="SRT generation complete!") | |
| return srt_content, message | |
| else: | |
| return "", message | |
| except Exception as e: | |
| return "", f"Error generating SRT: {str(e)}" | |
| def create_styled_srt(self, srt_content, font_size=24, font_color="white", | |
| outline_color="black", outline_width=1): | |
| """Create a styled SRT file with ASS-style formatting""" | |
| lines = srt_content.strip().split('\n') | |
| styled_lines = [] | |
| i = 0 | |
| while i < len(lines): | |
| if lines[i].strip().isdigit(): # Subtitle number | |
| styled_lines.append(lines[i]) | |
| i += 1 | |
| if i < len(lines): # Timestamp | |
| styled_lines.append(lines[i]) | |
| i += 1 | |
| # Collect all text lines for this subtitle | |
| text_lines = [] | |
| while i < len(lines) and lines[i].strip() != "": | |
| text_lines.append(lines[i]) | |
| i += 1 | |
| # Apply styling to text | |
| if text_lines: | |
| styled_text = ' '.join(text_lines) | |
| # Add basic styling tags | |
| styled_text = f"<font size='{font_size}' color='{font_color}'>{styled_text}</font>" | |
| styled_lines.append(styled_text) | |
| styled_lines.append("") # Empty line separator | |
| else: | |
| i += 1 | |
| return '\n'.join(styled_lines) | |
| def get_video_info(self, video_path): | |
| """Get video information using ffprobe""" | |
| try: | |
| cmd = [ | |
| 'ffprobe', '-v', 'quiet', '-print_format', 'json', | |
| '-show_format', '-show_streams', video_path | |
| ] | |
| result = subprocess.run(cmd, capture_output=True, text=True, check=True) | |
| info = json.loads(result.stdout) | |
| # Find video stream | |
| video_stream = None | |
| for stream in info['streams']: | |
| if stream['codec_type'] == 'video': | |
| video_stream = stream | |
| break | |
| if video_stream: | |
| width = int(video_stream['width']) | |
| height = int(video_stream['height']) | |
| duration = float(video_stream.get('duration', 0)) | |
| return width, height, duration | |
| except Exception as e: | |
| print(f"Error getting video info: {e}") | |
| return 1920, 1080, 0 # Default values | |
| def burn_subtitles(self, video_file, srt_content, font_size=24, font_color="white", | |
| position="bottom_center", outline_color="black", outline_width=1, | |
| progress=gr.Progress()): | |
| """Burn subtitles into video using FFmpeg""" | |
| if not video_file or not srt_content.strip(): | |
| return None, "Please provide both video file and SRT content" | |
| try: | |
| progress(0.1, desc="Preparing files...") | |
| # Create temporary SRT file | |
| srt_path = os.path.join(self.temp_dir, "subtitles.srt") | |
| styled_srt = self.create_styled_srt(srt_content, font_size, font_color, | |
| outline_color, outline_width) | |
| with open(srt_path, 'w', encoding='utf-8') as f: | |
| f.write(styled_srt) | |
| progress(0.2, desc="Getting video information...") | |
| # Get video info | |
| width, height, duration = self.get_video_info(video_file) | |
| progress(0.3, desc="Starting subtitle burning...") | |
| # Output file | |
| output_filename = f"output_with_subtitles_{font_size}px.mp4" | |
| output_path = os.path.join(self.temp_dir, output_filename) | |
| # Build FFmpeg command with subtitle filter | |
| cmd = [ | |
| 'ffmpeg', '-y', # Overwrite output files | |
| '-i', video_file, | |
| '-vf', f"""subtitles='{srt_path}':force_style='FontSize={font_size},PrimaryColour=&H{self.color_to_bgr_hex(font_color)},OutlineColour=&H{self.color_to_bgr_hex(outline_color)},Outline={outline_width},Alignment=2'""", | |
| '-c:a', 'copy', # Copy audio without re-encoding | |
| '-c:v', 'libx264', # Video codec | |
| '-preset', 'medium', # Encoding preset | |
| '-crf', '23', # Quality setting | |
| output_path | |
| ] | |
| progress(0.4, desc="Processing video (this may take a while)...") | |
| # Run FFmpeg | |
| process = subprocess.Popen( | |
| cmd, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| universal_newlines=True | |
| ) | |
| # Monitor progress | |
| while True: | |
| output = process.stderr.readline() | |
| if output == '' and process.poll() is not None: | |
| break | |
| if output and 'time=' in output: | |
| # Try to extract time for progress | |
| try: | |
| time_str = output.split('time=')[1].split()[0] | |
| time_parts = time_str.split(':') | |
| current_seconds = (float(time_parts[0]) * 3600 + | |
| float(time_parts[1]) * 60 + | |
| float(time_parts[2])) | |
| if duration > 0: | |
| prog = 0.4 + (current_seconds / duration) * 0.5 | |
| progress(min(prog, 0.9), desc=f"Processing: {time_str}") | |
| except: | |
| pass | |
| progress(0.95, desc="Finalizing...") | |
| return_code = process.poll() | |
| if return_code == 0: | |
| progress(1.0, desc="Complete!") | |
| return output_path, "Video processed successfully!" | |
| else: | |
| error_output = process.stderr.read() | |
| return None, f"FFmpeg error: {error_output}" | |
| except Exception as e: | |
| return None, f"Error processing video: {str(e)}" | |
| def color_to_bgr_hex(self, color): | |
| """Convert color name to BGR hex for FFmpeg""" | |
| color_map = { | |
| 'white': 'FFFFFF', | |
| 'black': '000000', | |
| 'red': '0000FF', | |
| 'green': '00FF00', | |
| 'blue': 'FF0000', | |
| 'yellow': '00FFFF', | |
| 'cyan': 'FFFF00', | |
| 'magenta': 'FF00FF', | |
| 'orange': '0080FF', | |
| 'purple': '800080', | |
| 'pink': 'FFB6C1', | |
| 'gray': '808080', | |
| 'grey': '808080' | |
| } | |
| return color_map.get(color.lower(), 'FFFFFF') | |
| def preview_subtitles(self, srt_content, font_size, font_color, position): | |
| """Generate a preview of how subtitles will look""" | |
| if not srt_content.strip(): | |
| return "No SRT content provided" | |
| lines = srt_content.strip().split('\n') | |
| preview_lines = [] | |
| # Extract first few subtitles for preview | |
| subtitle_count = 0 | |
| i = 0 | |
| while i < len(lines) and subtitle_count < 3: | |
| if lines[i].strip().isdigit(): | |
| subtitle_num = lines[i].strip() | |
| i += 1 | |
| if i < len(lines): | |
| timestamp = lines[i].strip() | |
| i += 1 | |
| text_lines = [] | |
| while i < len(lines) and lines[i].strip() != "": | |
| text_lines.append(lines[i].strip()) | |
| i += 1 | |
| if text_lines: | |
| text = ' '.join(text_lines) | |
| preview_lines.append(f"#{subtitle_num} [{timestamp}]") | |
| preview_lines.append(f"Text: \"{text}\"") | |
| preview_lines.append(f"Style: {font_size}px {font_color} at {position}") | |
| preview_lines.append("---") | |
| subtitle_count += 1 | |
| else: | |
| i += 1 | |
| return '\n'.join(preview_lines) if preview_lines else "No valid subtitles found" | |
| # Initialize the subtitle burner | |
| burner = AutomatedSubtitleBurner() | |
| # Create Gradio interface | |
| def create_interface(): | |
| with gr.Blocks(title="Automated AI Subtitle Video Captions", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 🎬 Automated SRT Subtitle Video Burner") | |
| gr.Markdown("Upload a video and either auto-generate subtitles or paste your own SRT content!") | |
| if not NEMO_AVAILABLE: | |
| gr.Markdown("⚠️ **Note**: Auto-transcription is disabled. Install NeMo toolkit for automatic SRT generation.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 📁 Input") | |
| video_input = gr.File( | |
| label="Upload Video File", | |
| file_types=[".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv"], | |
| type="filepath" | |
| ) | |
| with gr.Row(): | |
| if NEMO_AVAILABLE and burner.model: | |
| auto_generate_btn = gr.Button("🤖 Auto-Generate SRT", variant="secondary") | |
| else: | |
| auto_generate_btn = gr.Button("🤖 Auto-Generate SRT (Disabled)", variant="secondary", interactive=False) | |
| srt_input = gr.Textbox( | |
| label="SRT Content (Auto-generated or Manual)", | |
| placeholder="SRT content will appear here after auto-generation, or paste your own...", | |
| lines=12, | |
| max_lines=20 | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 🎨 Subtitle Styling") | |
| font_size = gr.Slider( | |
| minimum=8, | |
| maximum=72, | |
| value=24, | |
| step=1, | |
| label="Font Size (px)" | |
| ) | |
| font_color = gr.Dropdown( | |
| choices=["white", "black", "red", "green", "blue", "yellow", | |
| "cyan", "magenta", "orange", "purple", "pink", "gray"], | |
| value="white", | |
| label="Font Color" | |
| ) | |
| position = gr.Dropdown( | |
| choices=["top_left", "top_center", "top_right", | |
| "center_left", "center", "center_right", | |
| "bottom_left", "bottom_center", "bottom_right"], | |
| value="bottom_center", | |
| label="Position" | |
| ) | |
| outline_color = gr.Dropdown( | |
| choices=["black", "white", "red", "green", "blue", "yellow", | |
| "cyan", "magenta", "orange", "purple", "pink", "gray"], | |
| value="black", | |
| label="Outline Color" | |
| ) | |
| outline_width = gr.Slider( | |
| minimum=0, | |
| maximum=5, | |
| value=1, | |
| step=1, | |
| label="Outline Width" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### 👁️ Preview") | |
| preview_output = gr.Textbox( | |
| label="Subtitle Preview", | |
| lines=8, | |
| interactive=False | |
| ) | |
| preview_btn = gr.Button("🔍 Preview Subtitles", variant="secondary") | |
| with gr.Row(): | |
| process_btn = gr.Button("🔥 Burn Subtitles to Video", variant="primary", size="lg") | |
| with gr.Row(): | |
| with gr.Column(): | |
| output_video = gr.File(label="Download Processed Video") | |
| status_output = gr.Textbox(label="Status", interactive=False) | |
| # Event handlers | |
| if NEMO_AVAILABLE and burner.model: | |
| auto_generate_btn.click( | |
| fn=burner.auto_generate_srt, | |
| inputs=[video_input], | |
| outputs=[srt_input, status_output], | |
| show_progress=True | |
| ) | |
| preview_btn.click( | |
| fn=burner.preview_subtitles, | |
| inputs=[srt_input, font_size, font_color, position], | |
| outputs=preview_output | |
| ) | |
| process_btn.click( | |
| fn=burner.burn_subtitles, | |
| inputs=[video_input, srt_input, font_size, font_color, position, | |
| outline_color, outline_width], | |
| outputs=[output_video, status_output], | |
| show_progress=True | |
| ) | |
| # Auto-preview when inputs change | |
| for input_component in [srt_input, font_size, font_color, position]: | |
| input_component.change( | |
| fn=burner.preview_subtitles, | |
| inputs=[srt_input, font_size, font_color, position], | |
| outputs=preview_output | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| # Check if FFmpeg is available | |
| try: | |
| subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True) | |
| print("✅ FFmpeg found!") | |
| except (subprocess.CalledProcessError, FileNotFoundError): | |
| print("❌ FFmpeg not found! Please install FFmpeg and make sure it's in your PATH.") | |
| print("Download from: https://ffmpeg.org/download.html") | |
| exit(1) | |
| # Check transcription capability | |
| if NEMO_AVAILABLE and burner.model: | |
| print("✅ Auto-transcription enabled!") | |
| else: | |
| print("⚠️ Auto-transcription disabled. Install NeMo toolkit for automatic SRT generation:") | |
| print("pip install nemo_toolkit[asr]") | |
| # Launch the interface | |
| demo = create_interface() | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=True, | |
| debug=True | |
| ) | |