app.py

import gradio as gr
import os
import time
import tempfile
import requests
from urllib.parse import urlparse
from pydub import AudioSegment
import logging
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import yt_dlp

logging.basicConfig(level=logging.INFO)

sys.path.append("./faster-whisper")
from faster_whisper import WhisperModel, BatchedInferencePipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"

def download_audio(url, method_choice):
    parsed_url = urlparse(url)
    if parsed_url.netloc in ['www.youtube.com', 'youtu.be', 'youtube.com']:
        return download_youtube_audio(url, method_choice)
    else:
        return download_direct_audio(url, method_choice)

def download_youtube_audio(url, method_choice):
    methods = {
        'yt-dlp': youtube_dl_method,
        'pytube': pytube_method,
        'youtube-dl': youtube_dl_classic_method,
        'yt-dlp-alt': youtube_dl_alternative_method,
        'ffmpeg': ffmpeg_method,
        'aria2': aria2_method
    }
    method = methods.get(method_choice, youtube_dl_method)
    try:
        return method(url)
    except Exception as e:
        logging.error(f"Error downloading using {method_choice}: {str(e)}")
        return None

def youtube_dl_method(url):
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        return f"{info['id']}.mp3"

def pytube_method(url):
    from pytube import YouTube
    yt = YouTube(url)
    audio_stream = yt.streams.filter(only_audio=True).first()
    out_file = audio_stream.download()
    base, ext = os.path.splitext(out_file)
    new_file = base + '.mp3'
    os.rename(out_file, new_file)
    return new_file

def youtube_dl_classic_method(url):
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        return f"{info['id']}.mp3"

def youtube_dl_alternative_method(url):
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio', 
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
        'no_warnings': True,
        'quiet': True,
        'no_check_certificate': True,
        'prefer_insecure': True,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True) 
        return f"{info['id']}.mp3"

def ffmpeg_method(url):
    output_file = tempfile.mktemp(suffix='.mp3')
    command = ['ffmpeg', '-i', url, '-vn', '-acodec', 'libmp3lame', '-q:a', '2', output_file]
    subprocess.run(command, check=True, capture_output=True)
    return output_file

def aria2_method(url):  
    output_file = tempfile.mktemp(suffix='.mp3')
    command = ['aria2c', '--split=4', '--max-connection-per-server=4', '--out', output_file, url]
    subprocess.run(command, check=True, capture_output=True)
    return output_file

def download_direct_audio(url, method_choice):
    if method_choice == 'wget':
        return wget_method(url)
    else:
        try:
            response = requests.get(url)
            if response.status_code == 200:
                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
                    temp_file.write(response.content)
                    return temp_file.name
            else:
                raise Exception(f"Failed to download audio from {url}")
        except Exception as e:
            logging.error(f"Error downloading direct audio: {str(e)}")
            return None

def wget_method(url):
    output_file = tempfile.mktemp(suffix='.mp3')  
    command = ['wget', '-O', output_file, url]
    subprocess.run(command, check=True, capture_output=True)
    return output_file

def trim_audio(audio_path, start_time, end_time):
    audio = AudioSegment.from_file(audio_path)
    trimmed_audio = audio[start_time*1000:end_time*1000] if end_time else audio[start_time*1000:]
    trimmed_audio_path = tempfile.mktemp(suffix='.wav')
    trimmed_audio.export(trimmed_audio_path, format="wav")
    return trimmed_audio_path

def save_transcription(transcription):
    file_path = tempfile.mktemp(suffix='.txt')
    with open(file_path, 'w') as f:
        f.write(transcription)
    return file_path

def transcribe_audio(input_source, batch_size, download_method, start_time=None, end_time=None, verbose=False):
    try:
        model = WhisperModel("cstr/whisper-large-v3-turbo-int8_float32", device="auto", compute_type="int8")
        batched_model = BatchedInferencePipeline(model=model)

        if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')):
            audio_path = download_audio(input_source, download_method)
            if audio_path.startswith("Error"):
                yield f"Error: {audio_path}", "", None
                return
        else:
            audio_path = input_source

        if start_time is not None or end_time is not None:
            trimmed_audio_path = trim_audio(audio_path, start_time or 0, end_time)
            audio_path = trimmed_audio_path

        start_time_perf = time.time()
        segments, info = batched_model.transcribe(audio_path, batch_size=batch_size, initial_prompt=None)
        end_time_perf = time.time()

        transcription_time = end_time_perf - start_time_perf
        real_time_factor = info.duration / transcription_time
        audio_file_size = os.path.getsize(audio_path) / (1024 * 1024)

        metrics_output = (
            f"Language: {info.language}, Probability: {info.language_probability:.2f}\n"
            f"Duration: {info.duration:.2f}s, Duration after VAD: {info.duration_after_vad:.2f}s\n"  
            f"Transcription time: {transcription_time:.2f} seconds\n"
            f"Real-time factor: {real_time_factor:.2f}x\n"
            f"Audio file size: {audio_file_size:.2f} MB\n"
        )

        if verbose:
            yield metrics_output, "", None

        transcription = ""

        for segment in segments:
            transcription_segment = f"[{segment.start:.2f}s -> {segment.end:.2f}s] {segment.text}\n"
            transcription += transcription_segment

            if verbose:
                yield metrics_output, transcription, None

        transcription_file = save_transcription(transcription)
        yield metrics_output, transcription, transcription_file

    except Exception as e:
        yield f"An error occurred: {str(e)}", "", None

    finally:
        if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')):
            try:
                os.remove(audio_path)
            except:
                pass
        if start_time is not None or end_time is not None:
            try:
                os.remove(trimmed_audio_path)
            except:
                pass

iface = gr.Interface(
    fn=transcribe_audio,
    inputs=[
        gr.Textbox(label="Audio Source (Upload, URL, or YouTube URL)"),
        gr.Slider(minimum=1, maximum=32, step=1, value=16, label="Batch Size"),
        gr.Dropdown(choices=["yt-dlp", "pytube", "youtube-dl", "yt-dlp-alt", "ffmpeg", "aria2", "wget"], label="Download Method", value="yt-dlp"),
        gr.Number(label="Start Time (seconds)", value=0),  
        gr.Number(label="End Time (seconds)", value=0),
        gr.Checkbox(label="Verbose Output", value=False)
    ],
    outputs=[
        gr.Textbox(label="Transcription Metrics and Verbose Messages", lines=10),
        gr.Textbox(label="Transcription", lines=10),
        gr.File(label="Download Transcription")
    ],
    title="Multi-Model Transcription",
    description="Transcribe audio using with Whisper.",
    examples=[
        ["https://www.youtube.com/watch?v=daQ_hqA6HDo", 16, "yt-dlp", 0, None, False], 
        ["https://mcdn.podbean.com/mf/web/dir5wty678b6g4vg/HoP_453_-_The_Price_is_Right_-_Law_and_Economics_in_the_Second_Scholastic5yxzh.mp3", 16, "ffmpeg", 0, 300, True],
        ["path/to/local/audio.mp3", 16, "yt-dlp", 60, 180, False]
    ],
    cache_examples=False,
    live=True
)

iface.launch()