AutoSubGen / app.py
Athspi's picture
Update app.py
1d3d329 verified
raw
history blame
14.6 kB
import os
import google.generativeai as genai
from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, TextClip
from moviepy.video.tools.subtitles import SubtitlesClip
import tempfile
import logging
import srt
import datetime
from gtts import gTTS
from flask import Flask, request, render_template, send_from_directory, url_for, flash, session, redirect
from werkzeug.utils import secure_filename
import uuid
# --- Configuration ---
# Suppress moviepy logs which can be verbose
logging.basicConfig(level=logging.INFO)
logging.getLogger("moviepy").setLevel(logging.ERROR)
# Configure Gemini API
# IMPORTANT: Set your GEMINI_API_KEY as an environment variable
try:
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
except KeyError:
raise Exception("GEMINI_API_KEY environment variable not set. Please set it before running the app.")
# --- Flask App Initialization ---
app = Flask(__name__)
app.config['SECRET_KEY'] = os.urandom(24)
# Create a permanent directory for uploads if it doesn't exist
UPLOAD_FOLDER = os.path.join(os.getcwd(), 'user_uploads')
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
# --- Model and Language Configuration ---
# Create the Gemini model
generation_config = {
"temperature": 0.7,
"top_p": 0.9,
"top_k": 40,
"max_output_tokens": 8192,
"response_mime_type": "text/plain",
}
model = genai.GenerativeModel(
model_name="gemini-1.5-pro-latest",
generation_config=generation_config,
)
# List of all supported languages for the UI
SUPPORTED_LANGUAGES = [
"Auto Detect", "English", "Chinese", "German", "Spanish", "Russian", "Korean",
"French", "Japanese", "Portuguese", "Turkish", "Polish", "Catalan", "Dutch",
"Arabic", "Swedish", "Italian", "Indonesian", "Hindi", "Finnish", "Vietnamese",
"Hebrew", "Ukrainian", "Greek", "Malay", "Czech", "Romanian", "Danish",
"Hungarian", "Tamil", "Norwegian", "Thai", "Urdu", "Croatian", "Bulgarian",
"Lithuanian", "Latin", "Maori", "Malayalam", "Welsh", "Slovak", "Telugu",
"Persian", "Latvian", "Bengali", "Serbian", "Azerbaijani", "Slovenian",
"Kannada", "Estonian", "Macedonian", "Breton", "Basque", "Icelandic",
"Armenian", "Nepali", "Mongolian", "Bosnian", "Kazakh", "Albanian",
"Swahili", "Galician", "Marathi", "Punjabi", "Sinhala", "Khmer", "Shona",
"Yoruba", "Somali", "Afrikaans", "Occitan", "Georgian", "Belarusian",
"Tajik", "Sindhi", "Gujarati", "Amharic", "Yiddish", "Lao", "Uzbek",
"Faroese", "Haitian Creole", "Pashto", "Turkmen", "Nynorsk", "Maltese",
"Sanskrit", "Luxembourgish", "Burmese", "Tibetan", "Tagalog", "Malagasy",
"Assamese", "Tatar", "Hawaiian", "Lingala", "Hausa", "Bashkir", "Javanese",
"Sundanese"
]
# Language code mapping for Google Text-to-Speech (gTTS)
LANGUAGE_CODES = {
"English": "en", "Chinese": "zh", "German": "de", "Spanish": "es", "Russian": "ru",
"Korean": "ko", "French": "fr", "Japanese": "ja", "Portuguese": "pt", "Turkish": "tr",
"Polish": "pl", "Catalan": "ca", "Dutch": "nl", "Arabic": "ar", "Swedish": "sv",
"Italian": "it", "Indonesian": "id", "Hindi": "hi", "Finnish": "fi", "Vietnamese": "vi",
"Hebrew": "iw", "Ukrainian": "uk", "Greek": "el", "Malay": "ms", "Czech": "cs",
"Romanian": "ro", "Danish": "da", "Hungarian": "hu", "Tamil": "ta", "Norwegian": "no",
"Thai": "th", "Urdu": "ur", "Croatian": "hr", "Bulgarian": "bg", "Lithuanian": "lt",
"Latin": "la", "Malayalam": "ml", "Welsh": "cy", "Slovak": "sk", "Telugu": "te",
"Persian": "fa", "Latvian": "lv", "Bengali": "bn", "Serbian": "sr", "Slovenian": "sl",
"Kannada": "kn", "Estonian": "et", "Macedonian": "mk", "Armenian": "hy", "Nepali": "ne",
"Mongolian": "mn", "Bosnian": "bs", "Kazakh": "kk", "Albanian": "sq", "Swahili": "sw",
"Gujarati": "gu", "Lao": "lo", "Uzbek": "uz", "Maltese": "mt", "Burmese": "my",
"Tagalog": "tl", "Javanese": "jw", "Sundanese": "su", "Afrikaans": "af"
}
# --- Core Processing Functions ---
def extract_audio_from_video(video_path, audio_path):
"""Extract audio from a video file into a WAV file."""
try:
with VideoFileClip(video_path) as video:
video.audio.write_audiofile(audio_path, fps=16000, logger=None)
return audio_path
except Exception as e:
logging.error(f"Error extracting audio: {e}")
return None
def transcribe_audio_with_gemini(audio_path, source_language):
"""Transcribe audio using Gemini, requesting output in SRT format."""
try:
logging.info(f"Uploading audio file: {audio_path}")
audio_file = genai.upload_file(path=audio_path)
language_prompt = f"in {source_language}" if source_language != "Auto Detect" else ""
prompt = f"""
You are a professional transcriber. Transcribe this audio accurately and verbatim {language_prompt}.
Your output MUST be in the SRT (SubRip) format.
Example:
1
00:00:01,234 --> 00:00:05,678
This is the first sentence.
2
00:00:06,123 --> 00:00:09,456
This is the second sentence.
Ensure timestamps are precise. Respond ONLY with the transcription in the SRT format.
Do not add explanations, notes, or any other text outside of the valid SRT content.
"""
logging.info("Sending transcription request to Gemini...")
response = model.generate_content([prompt, audio_file], request_options={"timeout": 600})
genai.delete_file(audio_file.name) # Clean up the uploaded file
logging.info("Transcription received from Gemini.")
return response.text.strip()
except Exception as e:
logging.error(f"Error during Gemini transcription: {e}")
return None
def translate_srt(srt_text, target_language):
"""Translate an SRT file using Gemini while preserving its structure."""
try:
prompt = f"""
Translate the following SRT subtitles into {target_language}.
Preserve the SRT format perfectly (index numbers, timestamps, and structure).
Translate only the subtitle text itself.
Your output must be only the translated and valid SRT file content.
Do not add any explanations or extra text.
SRT Content to Translate:
{srt_text}
"""
logging.info(f"Sending translation request to Gemini for {target_language}...")
response = model.generate_content(prompt)
logging.info("Translation received.")
return response.text.strip()
except Exception as e:
logging.error(f"Error during translation: {e}")
return None
def generate_tts_audio(srt_text, language, tts_audio_path):
"""Generate a single TTS audio file from the text content of an SRT file."""
try:
subtitles = list(srt.parse(srt_text))
full_text = " ".join([sub.content.replace('\n', ' ') for sub in subtitles])
if not full_text:
return None
lang_code = LANGUAGE_CODES.get(language)
if not lang_code:
flash(f"Language '{language}' not supported for TTS, defaulting to English.", "warning")
lang_code = "en"
logging.info(f"Generating TTS audio in '{language}' ({lang_code})...")
tts = gTTS(text=full_text, lang=lang_code, slow=False)
tts.save(tts_audio_path)
logging.info(f"TTS audio saved to {tts_audio_path}")
return tts_audio_path
except Exception as e:
logging.error(f"Error generating TTS audio: {e}")
return None
def create_final_video(original_video_path, srt_text, new_audio_path, output_path):
"""
Creates the final video.
- If srt_text is provided, subtitles are burned in.
- If new_audio_path is provided, the original audio is replaced.
"""
try:
logging.info("Creating final video...")
original_clip = VideoFileClip(original_video_path)
# If new audio is provided, replace the original audio track
if new_audio_path:
tts_audio_clip = AudioFileClip(new_audio_path)
final_clip = original_clip.set_audio(tts_audio_clip)
else:
final_clip = original_clip
# If subtitle text is provided, burn it into the video
if srt_text:
with tempfile.NamedTemporaryFile(mode='w', suffix='.srt', delete=False, encoding='utf-8') as temp_srt:
temp_srt.write(srt_text)
srt_filepath = temp_srt.name
# Subtitle styling
generator = lambda txt: TextClip(
txt, font='Arial-Bold', fontsize=24, color='white',
stroke_color='black', stroke_width=1.5, method='caption',
size=(final_clip.w * 0.8, None) # Subtitles take 80% of video width
)
subtitles_clip = SubtitlesClip(srt_filepath, generator)
# Composite the video with the subtitles
final_video = CompositeVideoClip([final_clip, subtitles_clip.set_position(('center', 'bottom'))])
os.remove(srt_filepath) # Clean up temp srt file
else:
final_video = final_clip
# Write the final video file
final_video.write_videofile(output_path, codec='libx264', audio_codec='aac', threads=4, logger=None)
# Close all clips to release memory
if 'tts_audio_clip' in locals():
tts_audio_clip.close()
final_video.close()
logging.info(f"Final video saved to {output_path}")
return output_path
except Exception as e:
logging.error(f"Error creating final video: {e}")
return None
# --- Flask Routes ---
@app.route('/')
def index():
"""Render the main page. Clear session for a fresh start."""
session.clear()
return render_template('index.html', supported_languages=SUPPORTED_LANGUAGES)
@app.route('/process', methods=['POST'])
def process():
"""Handle the video processing request from the form."""
if 'video' not in request.files or request.files['video'].filename == '':
flash('No video file selected. Please upload a video.', 'error')
return redirect(url_for('index'))
video_file = request.files['video']
# --- Get form options ---
source_language = request.form.get('source_language', 'Auto Detect')
translate_to = request.form.get('translate_to', 'None')
add_tts = 'add_tts' in request.form
add_subtitles = 'add_subtitles' in request.form
# --- Setup a unique session directory for this request ---
session_id = str(uuid.uuid4())
session['session_id'] = session_id
session_dir = os.path.join(app.config['UPLOAD_FOLDER'], session_id)
os.makedirs(session_dir, exist_ok=True)
filename = secure_filename(video_file.filename)
original_video_path = os.path.join(session_dir, filename)
video_file.save(original_video_path)
results = {}
# 1. Extract Audio
audio_path = os.path.join(session_dir, "extracted_audio.wav")
if not extract_audio_from_video(original_video_path, audio_path):
flash('Failed to extract audio from the video. The file might be corrupted or in an unsupported format.', 'error')
return redirect(url_for('index'))
# 2. Transcribe Audio
original_srt_text = transcribe_audio_with_gemini(audio_path, source_language)
os.remove(audio_path) # Clean up extracted audio immediately
if not original_srt_text:
flash('Failed to transcribe the audio. The API call might have failed or the audio is silent.', 'error')
return redirect(url_for('index'))
original_srt_path = os.path.join(session_dir, "original_subtitles.srt")
with open(original_srt_path, "w", encoding="utf-8") as f:
f.write(original_srt_text)
results['original_srt_file'] = "original_subtitles.srt"
srt_for_final_video = None
tts_audio_path = None
final_srt_text = original_srt_text
# 3. Translate Subtitles (if requested)
if translate_to != "None":
translated_srt_text = translate_srt(original_srt_text, translate_to)
if translated_srt_text:
translated_srt_path = os.path.join(session_dir, "translated_subtitles.srt")
with open(translated_srt_path, "w", encoding="utf-8") as f:
f.write(translated_srt_text)
results['translated_srt_file'] = "translated_subtitles.srt"
final_srt_text = translated_srt_text # Use translated text for subsequent steps
else:
flash(f'Failed to translate subtitles to {translate_to}. Using original subtitles.', 'warning')
# 4. Generate TTS Audio (if requested)
if add_tts:
tts_lang = translate_to if translate_to != "None" else source_language
if tts_lang == 'Auto Detect':
flash('TTS language cannot be "Auto Detect". Defaulting to English.', 'warning')
tts_lang = 'English'
tts_audio_path_out = os.path.join(session_dir, "tts_audio.mp3")
tts_audio_path = generate_tts_audio(final_srt_text, tts_lang, tts_audio_path_out)
if tts_audio_path:
results['tts_audio_file'] = "tts_audio.mp3"
else:
flash('Failed to generate Text-to-Speech audio.', 'warning')
tts_audio_path = None # Ensure it's None if generation failed
# 5. Determine which subtitles to burn into the video (if requested)
if add_subtitles:
srt_for_final_video = final_srt_text
# 6. Create the final output video if any processing was requested
if add_subtitles or add_tts:
output_video_path = os.path.join(session_dir, "final_video.mp4")
if create_final_video(original_video_path, srt_for_final_video, tts_audio_path, output_video_path):
results['output_video_file'] = "final_video.mp4"
else:
flash('Failed to create the final processed video.', 'error')
return render_template('index.html',
supported_languages=SUPPORTED_LANGUAGES,
results=results,
session_id=session_id)
@app.route('/download/<session_id>/<path:filename>')
def download_file(session_id, filename):
"""Serve files from the session-specific directory for download."""
directory = os.path.join(app.config['UPLOAD_FOLDER'], session_id)
return send_from_directory(directory, filename, as_attachment=True)
# --- Run the App ---
if __name__ == '__main__':
app.run(debug=True, port=5000)