Create app.py

app.py

@@ -0,0 +1,223 @@
+import whisper
+import yt_dlp
+import moviepy.editor as mp
+import os
+import torch
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import pipeline
+import srt
+from datetime import timedelta
+import gradio as gr
+import torchaudio
+import whisper.tokenizer
+
+# -----------------------------
+# Helper Functions
+# -----------------------------
+
+def download_youtube_audio(url):
+    ydl_opts = {
+        'format': 'bestaudio/best',
+        'outtmpl': 'audio.%(ext)s',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'mp3',
+            'preferredquality': '192',
+        }],
+    }
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info = ydl.extract_info(url, download=True)
+            return "audio.mp3"
+    except Exception as e:
+        raise RuntimeError(f"Error downloading audio: {str(e)}")
+
+def extract_audio_from_video(video_path):
+    try:
+        clip = mp.VideoFileClip(video_path)
+        clip.audio.write_audiofile("audio.mp3")
+        return "audio.mp3"
+    except Exception as e:
+        raise RuntimeError(f"Error extracting audio: {str(e)}")
+
+def generate_srt(segments):
+    subs = []
+    for i, seg in enumerate(segments):
+        start = timedelta(seconds=seg['start'])
+        end = timedelta(seconds=seg['end'])
+        text = seg['text'].strip()
+        if text:
+            subs.append(srt.Subtitle(index=i+1, start=start, end=end, content=text))
+    return srt.compose(subs)
+
+# -----------------------------
+# Transcription Functions
+# -----------------------------
+
+def transcribe_kotani(audio_path):
+    model = whisper.load_model("small", download_root=".")
+    result = model.transcribe(audio_path, language=None)  # auto-detect
+    return result["segments"], result["language"]
+
+def transcribe_khaiii(audio_path):
+    processor = Wav2Vec2Processor.from_pretrained("khaiii/wav2vec2-xls1r-aishell-korean")
+    model = Wav2Vec2ForCTC.from_pretrained("khaiii/wav2vec2-xls1r-aishell-korean")
+    speech, sr = torchaudio.load(audio_path)
+    input_values = processor(speech.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
+    logits = model(input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = processor.batch_decode(predicted_ids)[0]
+    duration = len(speech) / sr
+    return [{"start": 0, "end": duration, "text": transcription}], "ko"
+
+# -----------------------------
+# Translation Function
+# -----------------------------
+
+def translate_text(text, src_lang, tgt_lang="en"):
+    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
+    try:
+        translator = pipeline("translation", model=model_name)
+        translated = translator(text, max_length=400)
+        return translated[0]['translation_text']
+    except Exception as e:
+        return f"[Translation error: {str(e)}]"
+
+# -----------------------------
+# Main Processing Function
+# -----------------------------
+
+def process_video(youtube_url, video_file, selected_model, translate, target_lang):
+    status = "⏳ Starting..."
+    yield status, "", None
+
+    try:
+        # Step 1: Extract audio
+        if youtube_url:
+            status = "📥 Downloading YouTube audio..."
+            yield status, "", None
+            audio_path = download_youtube_audio(youtube_url)
+        elif video_file:
+            status = "📼 Extracting audio from video..."
+            yield status, "", None
+            audio_path = extract_audio_from_video(video_file.name)
+        else:
+            yield "❌ Please provide a video or YouTube URL", "", None
+            return
+
+        # Step 2: Transcribe
+        if selected_model == "kotani":
+            status = "🎙️ Transcribing using Kotani Whisper Small..."
+            yield status, "", None
+            segments, lang = transcribe_kotani(audio_path)
+        else:
+            status = "🎙️ Transcribing using Khaiii Wav2Vec2..."
+            yield status, "", None
+            segments, lang = transcribe_khaiii(audio_path)
+
+        lang_desc = whisper.tokenizer.LLANGUAGES.get(lang, lang.upper())
+
+        # Step 3: Translate if needed
+        if translate:
+            status = f"🌐 Translating {lang_desc} to {target_lang.upper()}..."
+            yield status, "", None
+            translated_segments = []
+            for seg in segments:
+                translated = translate_text(seg['text'], lang, target_lang)
+                translated_segments.append({**seg, "text": translated})
+            segments = translated_segments
+
+        # Step 4: Generate SRT
+        status = "📝 Generating subtitle file..."
+        yield status, "", None
+        srt_content = generate_srt(segments)
+
+        with open("output.srt", "w") as f:
+            f.write(srt_content)
+
+        preview = srt_content[:1000] + ("\n..." if len(srt_content) > 1000 else "")
+        status = f"✅ Done! ({lang_desc})"
+        yield status, preview, "output.srt"
+
+    except Exception as e:
+        yield f"❌ Error: {str(e)}", "", None
+
+# -----------------------------
+# UI Layout
+# -----------------------------
+
+model_desc_kotani = """
+<div style="border:1px solid #ddd; padding: 10px; border-radius:8px;">
+  <strong>Kotani Whisper Small</strong><br>
+  ▪ Fast & multilingual<br>
+  ▪ Good for quick subtitles<br>
+  ▪ Moderate accuracy for Korean
+</div>
+"""
+
+model_desc_khaiii = """
+<div style="border:1px solid #ddd; padding: 10px; border-radius:8px;">
+  <strong>Khaiii Wav2Vec2</strong><br>
+  ▪ Best Korean speech recognition<br>
+  ▪ Slower but highly accurate<br>
+  ▪ Only supports Korean
+</div>
+"""
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🌍 Multilingual Subtitle Generator")
+    gr.Markdown("Upload a video or paste a YouTube link. Automatically detect language and optionally translate subtitles.")
+
+    selected_model = gr.State(value="kotani")  # default model
+
+    gr.Markdown("### 🔍 Choose ASR Model")
+    with gr.Row():
+        with gr.Column():
+            kotani_btn = gr.Button("✅ Select Kotani Whisper Small")
+            gr.HTML(model_desc_kotani)
+        with gr.Column():
+            khaiii_btn = gr.Button("✅ Select Khaiii Wav2Vec2")
+            gr.HTML(model_desc_khaiii)
+
+    selected_model = gr.State(value="kotani")
+
+    def select_kotani():
+        return "kotani"
+
+    def select_khaiii():
+        return "khaiii"
+
+    kotani_btn.click(fn=select_kotani, outputs=selected_model)
+    khaiii_btn.click(fn=select_khaiii, outputs=selected_model)
+
+    gr.Markdown("### 📥 Input Source")
+    with gr.Row():
+        youtube_url = gr.Textbox(label="YouTube URL", scale=2)
+        video_upload = gr.File(label="Upload Video", type="file", scale=1)
+
+    gr.Markdown("### 🌍 Translation Options")
+    with gr.Row():
+        translate_checkbox = gr.Checkbox(label="Translate to another language?")
+        target_lang = gr.Textbox(label="Target Language Code (e.g., 'en')", value="en", visible=False)
+
+    def toggle_translate(checked):
+        return gr.update(visible=checked)
+
+    translate_checkbox.change(fn=toggle_translate, inputs=translate_checkbox, outputs=target_lang)
+
+    status_box = gr.Textbox(label="Status", interactive=False)
+    subtitle_preview = gr.Textbox(label="Generated Subtitles", lines=10)
+    download_file = gr.File(label="Download .srt File")
+
+    submit_btn = gr.Button("🎬 Generate Subtitles")
+
+    submit_btn.click(
+        fn=process_video,
+        inputs=[youtube_url, video_upload, selected_model, translate_checkbox, target_lang],
+        outputs=[status_box, subtitle_preview, download_file]
+    )
+
+demo.queue(concurrency_count=1, max_size=20)
+
+if __name__ == "__main__":
+    demo.launch()