Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import time | |
| import sys | |
| import tempfile | |
| import subprocess | |
| 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) | |
| # Clone and install faster-whisper from GitHub | |
| try: | |
| subprocess.run(["git", "clone", "https://github.com/SYSTRAN/faster-whisper.git"], check=True) | |
| subprocess.run(["pip", "install", "-e", "./faster-whisper"], check=True) | |
| except subprocess.CalledProcessError as e: | |
| logging.error(f"Error during faster-whisper installation: {e}") | |
| sys.exit(1) | |
| sys.path.append("./faster-whisper") | |
| from faster_whisper import WhisperModel | |
| from faster_whisper.transcribe import 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 get_model_options(pipeline_type): | |
| if pipeline_type == "faster-batched": | |
| return ["cstr/whisper-large-v3-turbo-int8_float32"] | |
| elif pipeline_type == "faster-sequenced": | |
| return ["deepdml/faster-whisper-large-v3-turbo-ct2"] | |
| elif pipeline_type == "transformers": | |
| return ["openai/whisper-large-v3"] | |
| return [] | |
| def transcribe_audio(input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time=None, end_time=None, verbose=False): | |
| try: | |
| logging.info(f"Transcription parameters: pipeline_type={pipeline_type}, model_id={model_id}, dtype={dtype}, batch_size={batch_size}, download_method={download_method}") | |
| verbose_messages = f"Starting transcription with parameters:\nPipeline Type: {pipeline_type}\nModel ID: {model_id}\nData Type: {dtype}\nBatch Size: {batch_size}\nDownload Method: {download_method}\n" | |
| if pipeline_type == "faster-batched": | |
| model = WhisperModel(model_id, device="auto", compute_type=dtype) | |
| pipeline = BatchedInferencePipeline(model=model) | |
| elif pipeline_type == "faster-sequenced": | |
| model = WhisperModel(model_id) | |
| pipeline = model.transcribe | |
| elif pipeline_type == "transformers": | |
| torch_dtype = torch.float16 if dtype == "float16" else torch.float32 | |
| model = AutoModelForSpeechSeq2Seq.from_pretrained( | |
| model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True | |
| ) | |
| model.to(device) | |
| processor = AutoProcessor.from_pretrained(model_id) | |
| pipeline = pipeline( | |
| "automatic-speech-recognition", | |
| model=model, | |
| tokenizer=processor.tokenizer, | |
| feature_extractor=processor.feature_extractor, | |
| chunk_length_s=30, | |
| batch_size=batch_size, | |
| return_timestamps=True, | |
| torch_dtype=torch_dtype, | |
| device=device, | |
| ) | |
| else: | |
| raise ValueError("Invalid pipeline type") | |
| if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')): | |
| audio_path = download_audio(input_source, download_method) | |
| verbose_messages += f"Audio file downloaded: {audio_path}\n" | |
| 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 | |
| verbose_messages += f"Audio trimmed from {start_time} to {end_time}\n" | |
| start_time_perf = time.time() | |
| if pipeline_type in ["faster-batched", "faster-sequenced"]: | |
| segments, info = pipeline(audio_path, batch_size=batch_size) | |
| else: | |
| result = pipeline(audio_path) | |
| segments = result["chunks"] | |
| end_time_perf = time.time() | |
| transcription_time = end_time_perf - start_time_perf | |
| audio_file_size = os.path.getsize(audio_path) / (1024 * 1024) | |
| metrics_output = ( | |
| f"Transcription time: {transcription_time:.2f} seconds\n" | |
| f"Audio file size: {audio_file_size:.2f} MB\n" | |
| ) | |
| if verbose: | |
| yield verbose_messages + metrics_output, "", None | |
| transcription = "" | |
| for segment in segments: | |
| transcription_segment = ( | |
| f"[{segment.start:.2f}s -> {segment.end:.2f}s] {segment.text}\n" | |
| if pipeline_type in ["faster-batched", "faster-sequenced"] else | |
| f"[{segment['timestamp'][0]:.2f}s -> {segment['timestamp'][1]:.2f}s] {segment['text']}\n" | |
| ) | |
| transcription += transcription_segment | |
| if verbose: | |
| yield verbose_messages + metrics_output, transcription, None | |
| transcription_file = save_transcription(transcription) | |
| yield verbose_messages + metrics_output, transcription, transcription_file | |
| except Exception as e: | |
| logging.error(f"An error occurred during transcription: {str(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 | |
| def update_model_dropdown(pipeline_type): | |
| model_choices = get_model_options(pipeline_type) | |
| return gr.Dropdown.update(choices=model_choices, value=model_choices[0]) | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Multi-Pipeline Transcription") | |
| gr.Markdown("Transcribe audio using multiple pipelines and models.") | |
| with gr.Row(): | |
| input_source = gr.Textbox(label="Audio Source (Upload, URL, or YouTube URL)") | |
| pipeline_type = gr.Dropdown(choices=["faster-batched", "faster-sequenced", "transformers"], label="Pipeline Type", value="faster-batched") | |
| model_id = gr.Dropdown(label="Model", choices=get_model_options("faster-batched"), value=get_model_options("faster-batched")[0]) | |
| with gr.Row(): | |
| dtype = gr.Dropdown(choices=["int8", "float16", "float32"], label="Data Type", value="int8") | |
| batch_size = gr.Slider(minimum=1, maximum=32, step=1, value=16, label="Batch Size") | |
| download_method = gr.Dropdown(choices=["yt-dlp", "pytube", "youtube-dl", "yt-dlp-alt", "ffmpeg", "aria2", "wget"], label="Download Method", value="yt-dlp") | |
| with gr.Row(): | |
| start_time = gr.Number(label="Start Time (seconds)", value=0) | |
| end_time = gr.Number(label="End Time (seconds)", value=0) | |
| verbose = gr.Checkbox(label="Verbose Output", value=False) | |
| transcribe_button = gr.Button("Transcribe") | |
| with gr.Row(): | |
| metrics_output = gr.Textbox(label="Transcription Metrics and Verbose Messages", lines=10) | |
| transcription_output = gr.Textbox(label="Transcription", lines=10) | |
| transcription_file = gr.File(label="Download Transcription") | |
| pipeline_type.change(update_model_dropdown, inputs=pipeline_type, outputs=model_id) | |
| transcribe_button.click( | |
| transcribe_audio, | |
| inputs=[input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time, end_time, verbose], | |
| outputs=[metrics_output, transcription_output, transcription_file] | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["https://www.youtube.com/watch?v=daQ_hqA6HDo", "faster-batched", "cstr/whisper-large-v3-turbo-int8_float32", "int8", 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", "faster-sequenced", "deepdml/faster-whisper-large-v3-turbo-ct2", "float16", 1, "ffmpeg", 0, 300, True], | |
| ["path/to/local/audio.mp3", "transformers", "openai/whisper-large-v3", "float16", 16, "yt-dlp", 60, 180, False] | |
| ], | |
| inputs=[input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time, end_time, verbose], | |
| ) | |
| iface.launch() |