app.py

import uuid
import gradio as gr
import io
import os

from transformers import pipeline
import torch
import yt_dlp
from silero_vad import load_silero_vad, get_speech_timestamps
import numpy as np
import pydub
from litellm import completion

# --- Language List ---
LANGUAGES = ['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', 'myanmar', 'tibetan', 'tagalog', 'malagasy', 'assamese', 'tatar', 'hawaiian', 'lingala', 'hausa', 'bashkir', 'javanese', 'sundanese', 'cantonese', 'burmese', 'valencian', 'flemish', 'haitian', 'letzeburgesch', 'pushto', 'panjabi', 'moldavian', 'moldovan', 'sinhalese', 'castilian', 'mandarin']

# --- Model Loading and Caching ---
def load_transcriber(_device):
    """Loads the Whisper transcription model."""
    transcriber = pipeline(model="openai/whisper-large-v3-turbo", device=_device)
    return transcriber

def load_vad_model():
    """Loads the Silero VAD model."""
    return load_silero_vad()

# --- Audio Processing Functions ---
def download_and_convert_audio(video_url, audio_format="wav"):
    """Downloads and converts audio from a YouTube video.

    Args:
        video_url (str): The URL of the YouTube video.
        audio_format (str): The desired audio format (e.g., "wav", "mp3").

    Returns:
        tuple: (audio_bytes, audio_format, info_dict) or (None, None, None) on error.
    """
    try:
        ydl_opts = {
            'format': f'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': audio_format,
            }],
            'outtmpl': '%(id)s.%(ext)s',
            'noplaylist': True,
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(video_url, download=False)
            if 'entries' in info:
                info = info['entries'][0]
            video_id = info['id']
            filename = f"{video_id}.{audio_format}"

            audio_formats = [f for f in info.get('formats', []) if f.get('acodec') != 'none' and f.get('vcodec') == 'none']
            if not audio_formats:
                print(f"No audio-only format found. Downloading and converting from best video format to {audio_format}.")
                ydl_opts['format'] = 'best'

            ydl.download([video_url])
            print(f"Audio downloaded and converted to {audio_format}.")

            with open(filename, 'rb') as audio_file:
                audio_bytes = audio_file.read()

            os.remove(filename)
            return audio_bytes, audio_format, info
    except Exception as e:
        print(f"Error during download or conversion: {e}")
        return None, None, None

def split_audio_by_vad(audio_data: bytes, ext: str, _vad_model, sensitivity: float, max_duration: int = 30, return_seconds: bool = True):
    """Splits audio into chunks based on voice activity detection (VAD).

    Args:
        audio_data (bytes): The audio data as bytes.
        ext (str): The audio file extension.
        _vad_model: The VAD model.
        sensitivity (float): The VAD sensitivity (0.0 to 1.0).
        max_duration (int): The maximum duration of each chunk in seconds.
        return_seconds (bool): Whether to return timestamps in seconds.

    Returns:
        list: A list of dictionaries, where each dictionary represents an audio chunk.
              Returns an empty list if no speech segments are detected or an error occurs.
    """
    
    if not audio_data:
        print("No audio data received.")
        return []

    try:
        audio = pydub.AudioSegment.from_file(io.BytesIO(audio_data), format=ext)
        rate = audio.frame_rate
        
        # Convert to mono if stereo for compatibility with VAD
        if audio.channels > 1:
            audio = audio.set_channels(1)

        # Calculate dynamic VAD parameters based on sensitivity
        window_size_samples = int(512 + (1536 - 512) * (1 - sensitivity))
        speech_threshold = 0.5 + (0.95 - 0.5) * sensitivity
        
        samples = np.array(audio.get_array_of_samples())

        speech_timestamps = get_speech_timestamps(
            samples,
            _vad_model,
            sampling_rate=rate,
            return_seconds=return_seconds,
            window_size_samples=window_size_samples,
            threshold=speech_threshold,
        )

        if not speech_timestamps:
            print("No speech segments detected.")
            return []

        speech_timestamps[0]["start"] = 0.
        speech_timestamps[-1]['end'] = audio.duration_seconds
        for i, chunk in enumerate(speech_timestamps[1:], start=1):
            chunk["start"] = speech_timestamps[i - 1]['end']

        aggregated_segments = []
        if speech_timestamps:
            current_segment_start = speech_timestamps[0]['start']
            current_segment_end = speech_timestamps[0]['end']
            for segment in speech_timestamps[1:]:
                if segment['start'] - current_segment_start >= max_duration:
                    aggregated_segments.append({'start': current_segment_start, 'end': current_segment_end})
                    current_segment_start = segment['start']
                    current_segment_end = segment['end']
                else:
                    current_segment_end = segment['end']
            aggregated_segments.append({'start': current_segment_start, 'end': current_segment_end})
        
        if not aggregated_segments:
            return []

        chunks = []
        for segment in aggregated_segments:
            start_ms = int(segment['start'] * 1000)
            end_ms = int(segment['end'] * 1000)
            chunk = audio[start_ms:end_ms]
            chunk_io = io.BytesIO()
            chunk.export(chunk_io, format=ext)
            chunks.append({
                'data': chunk_io.getvalue(),
                'start': segment['start'],
                'end': segment['end']
            })
            chunk_io.close()
        return chunks
    except Exception as e:
        print(f"Error processing audio in split_audio_by_vad: {str(e)}")
        return []
    finally:
        if 'audio' in locals():
            del audio
        if 'samples' in locals():
            del samples

def transcribe_batch(batch, _transcriber, language=None):
    """Transcribes a batch of audio chunks.

    Args:
        batch (list): A list of audio chunk dictionaries.
        _transcriber: The transcription model.
        language (str, optional): The language of the audio (e.g., "en", "es"). Defaults to None (auto-detection).

    Returns:
        list: A list of dictionaries, each containing the transcription, start, and end time of a chunk.
              Returns an empty list if an error occurs.
    """
    transcriptions = []
    for i, chunk_data in enumerate(batch):
        try:
            generate_kwargs = {
                "task": "transcribe",
                "return_timestamps": True,
                "language": language
            }

            transcription = _transcriber(
                chunk_data['data'],
                generate_kwargs=generate_kwargs
            )
            transcriptions.append({
                'text': transcription["text"],
                'start': chunk_data['start'],
                'end': chunk_data['end']}
            )
        except Exception as e:
            print(f"Error transcribing chunk {i}: {str(e)}")
            return []
    return transcriptions

def format_seconds(seconds):
    """Formats seconds into HH:MM:SS string."""
    minutes, seconds = divmod(seconds, 60)
    hours, minutes = divmod(minutes, 60)
    return f"{int(hours):02}:{int(minutes):02}:{int(seconds):02}"

def download_video(video_url, video_format):
    """Downloads video from YouTube using yt-dlp."""
    try:
        ydl_opts = {
            'format': f'bestvideo[ext={video_format}]+bestaudio[ext=m4a]/best[ext={video_format}]/best',
            'outtmpl': '%(title)s.%(ext)s',
            'noplaylist': True,
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(video_url, download=True)
            video_filename = ydl.prepare_filename(info_dict)
            video_title = info_dict.get("title", "video")
            print(f"Video downloaded: {video_title}")

            with open(video_filename, 'rb') as video_file:
                video_bytes = video_file.read()

            os.remove(video_filename)

            return video_bytes, video_filename, info_dict
    except Exception as e:
        print(f"Error during video download: {e}")
        return None, None, None

def format_transcript(input_transcription):
    """Formats the transcription using the Gemini large language model."""

    os.environ["GEMINI_API_KEY"] = "AIzaSyBWmOE2XMVpiHuUI4YzgOVzUoENfDeXe8s"

    sys_prompt = """
    *   Format the provided video transcription as a polished piece of written text.
*   **The output must be in the same language as the input; do not translate it.**
*   Focus on clarity, readability, and consistency, adhering to the conventions of that specific language.
*   Restructure sentences for improved flow and correct grammatical errors.
*   **Edits should strictly enhance readability without altering the original meaning or nuances of the raw transcription.**
*   Italicize or quote any text that is read aloud, clearly distinguishing it from the surrounding explanations.
*   Eliminate unnecessary repetitions unless they are used for emphasis.
*   **Do not add any information not present in the original transcript.**
*   **Do not remove timestamps.**
*   **Output only the formatted transcription.**
    """.strip()
    messages = [{"content": sys_prompt, "role": "system"},
                 {"content": f"Format the following video transcription: {input_transcription}", "role": "user"}]

    response = completion(model="gemini/gemini-2.0-flash-exp", messages=messages)
    formatted_text = response.choices[0].message.content
    return formatted_text

def process_transcription(video_url, vad_sensitivity, batch_size, transcriber, vad_model, audio_format, language=None):
    """Downloads, processes, and transcribes the audio from a YouTube video.

    Args:
        video_url (str): The URL of the YouTube video.
        vad_sensitivity (float): The VAD sensitivity.
        batch_size (int): The batch size for transcription.
        transcriber: The transcription model.
        vad_model: The VAD model.
        language (str, optional): The language of the audio. Defaults to None.

    Returns:
        tuple: (full_transcription, audio_data, audio_format, info) or (None, None, None, None) on error.
    """
    audio_data, audio_format, info = download_and_convert_audio(video_url, audio_format)
    if not audio_data:
        return None, None, None, None

    chunks = split_audio_by_vad(audio_data, audio_format, vad_model, vad_sensitivity)
    if not chunks:
        return None, None, None, None

    total_chunks = len(chunks)
    transcriptions = []
    for i in range(0, total_chunks, batch_size):
        batch = chunks[i:i + batch_size]
        batch_transcriptions = transcribe_batch(batch, transcriber, language)
        transcriptions.extend(batch_transcriptions)

    full_transcription = ""
    for chunk in transcriptions:
        start_time = format_seconds(chunk['start'])
        end_time = format_seconds(chunk['end'])
        full_transcription += f"[{start_time} - {end_time}]: {chunk['text'].strip()}\n\n"

    return full_transcription, audio_data, audio_format, info

def main(video_url, language, batch_size, transcribe_option, download_audio_option, download_video_option, vad_sensitivity, audio_format, video_format, format_option):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    transcriber = load_transcriber(device)
    vad_model = load_vad_model()

    selected_language = language.lower() if language != "Auto-Detect" else None

    full_transcription = None
    formatted_transcription = None
    audio_data = None
    info = None
    video_data = None
    video_filename = None

    if transcribe_option:
        full_transcription, audio_data, audio_format, info = process_transcription(video_url, vad_sensitivity, batch_size, transcriber, vad_model, audio_format, selected_language)
        if full_transcription and format_option:
            formatted_transcription = format_transcript(full_transcription)

    if download_audio_option:
        if audio_data is None or audio_format is None:
            audio_data, audio_format, info = download_and_convert_audio(video_url, audio_format)

    if download_video_option:
        video_data, video_filename, info = download_video(video_url, video_format)

    return full_transcription, formatted_transcription, audio_data, audio_format, video_data, video_filename

iface = gr.Interface(
    fn=main,
    inputs=[
        gr.Textbox(label="YouTube Video Link"),
        gr.Dropdown(["Auto-Detect"] + LANGUAGES, label="Language", default="Auto-Detect"),
        gr.Number(label="Batch Size", value=2, precision=0),
        gr.Checkbox(label="Transcribe", value=True),
        gr.Checkbox(label="Download Audio", value=False),
        gr.Checkbox(label="Download Video", value=False),
        gr.Slider(label="VAD Sensitivity", minimum=0.0, maximum=1.0, value=0.1, step=0.05),
        gr.Dropdown(["wav", "mp3", "ogg", "flac"], label="Audio Format", default="wav"),
        gr.Dropdown(["mp4", "webm"], label="Video Format", default="mp4"),
        gr.Checkbox(label="Format Text", value=True)
    ],
    outputs=[
        gr.Textbox(label="Transcription"),
        gr.Textbox(label="Formatted Transcription"),
        gr.File(label="Audio File"),
        gr.Textbox(label="Audio Format"),
        gr.File(label="Video File"),
        gr.Textbox(label="Video Filename")
    ],
    title="YouTube Video Transcriber",
    description="This app allows you to transcribe YouTube videos and format the transcription using a large language model. You can also download the audio and video."
)

iface.launch()