import os
import torch
import whisper
import gradio as gr
import torchaudio
from abc import ABC, abstractmethod
from typing import BinaryIO, Union, Tuple, List
import numpy as np
from datetime import datetime
from faster_whisper.vad import VadOptions
from dataclasses import astuple
import gc
from copy import deepcopy
from modules.vad.silero_vad import merge_chunks, Segment
from modules.uvr.music_separator import MusicSeparator
from modules.utils.paths import (WHISPER_MODELS_DIR, DIARIZATION_MODELS_DIR, OUTPUT_DIR, DEFAULT_PARAMETERS_CONFIG_PATH,
UVR_MODELS_DIR)
from modules.utils.subtitle_manager import get_srt, get_vtt, get_txt, get_plaintext, get_csv, write_file, safe_filename
from modules.utils.youtube_manager import get_ytdata, get_ytaudio
from modules.utils.files_manager import get_media_files, format_gradio_files, load_yaml, save_yaml
from modules.whisper.whisper_parameter import *
from modules.diarize.diarizer import Diarizer
from modules.vad.silero_vad import SileroVAD
from modules.translation.nllb_inference import NLLBInference
from modules.translation.nllb_inference import NLLB_AVAILABLE_LANGS
import faster_whisper
class WhisperBase(ABC):
def __init__(self,
model_dir: str = WHISPER_MODELS_DIR,
diarization_model_dir: str = DIARIZATION_MODELS_DIR,
uvr_model_dir: str = UVR_MODELS_DIR,
output_dir: str = OUTPUT_DIR,
):
self.model_dir = model_dir
self.output_dir = output_dir
os.makedirs(self.output_dir, exist_ok=True)
os.makedirs(self.model_dir, exist_ok=True)
self.diarizer = Diarizer(
model_dir=diarization_model_dir
)
self.vad = SileroVAD()
self.music_separator = MusicSeparator(
model_dir=uvr_model_dir,
output_dir=os.path.join(output_dir, "UVR")
)
self.model = None
self.current_model_size = None
self.available_models = whisper.available_models()
self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
#self.translatable_models = ["large", "large-v1", "large-v2", "large-v3"]
self.translatable_models = whisper.available_models()
self.device = self.get_device()
self.available_compute_types = ["float16", "float32"]
self.current_compute_type = "float16" if self.device == "cuda" else "float32"
@abstractmethod
def transcribe(self,
audio: Union[str, BinaryIO, np.ndarray],
progress: gr.Progress = gr.Progress(),
*whisper_params,
):
"""Inference whisper model to transcribe"""
pass
@abstractmethod
def update_model(self,
model_size: str,
compute_type: str,
progress: gr.Progress = gr.Progress()
):
"""Initialize whisper model"""
pass
def run(self,
audio: Union[str, BinaryIO, np.ndarray],
progress: gr.Progress = gr.Progress(),
add_timestamp: bool = True,
*whisper_params,
) -> Tuple[List[dict], float]:
"""
Run transcription with conditional pre-processing and post-processing.
The VAD will be performed to remove noise from the audio input in pre-processing, if enabled.
The diarization will be performed in post-processing, if enabled.
Parameters
----------
audio: Union[str, BinaryIO, np.ndarray]
Audio input. This can be file path or binary type.
progress: gr.Progress
Indicator to show progress directly in gradio.
add_timestamp: bool
Whether to add a timestamp at the end of the filename.
*whisper_params: tuple
Parameters related with whisper. This will be dealt with "WhisperParameters" data class
Returns
----------
segments_result: List[dict]
list of dicts that includes start, end timestamps and transcribed text
elapsed_time: float
elapsed time for running
"""
start_time = datetime.now()
params = WhisperParameters.as_value(*whisper_params)
# Get the offload params
default_params = load_yaml(DEFAULT_PARAMETERS_CONFIG_PATH)
whisper_params = default_params["whisper"]
diarization_params = default_params["diarization"]
bool_whisper_enable_offload = whisper_params["enable_offload"]
bool_diarization_enable_offload = diarization_params["enable_offload"]
if params.lang is None:
pass
elif params.lang == "Automatic Detection":
params.lang = None
else:
language_code_dict = {value: key for key, value in whisper.tokenizer.LANGUAGES.items()}
params.lang = language_code_dict[params.lang]
if params.is_bgm_separate:
music, audio, _ = self.music_separator.separate(
audio=audio,
model_name=params.uvr_model_size,
device=params.uvr_device,
segment_size=params.uvr_segment_size,
save_file=params.uvr_save_file,
progress=progress
)
if audio.ndim >= 2:
audio = audio.mean(axis=1)
if self.music_separator.audio_info is None:
origin_sample_rate = 16000
else:
origin_sample_rate = self.music_separator.audio_info.sample_rate
audio = self.resample_audio(audio=audio, original_sample_rate=origin_sample_rate)
if params.uvr_enable_offload:
self.music_separator.offload()
elapsed_time_bgm_sep = datetime.now() - start_time
origin_audio = deepcopy(audio)
if params.vad_filter:
# Explicit value set for float('inf') from gr.Number()
if params.max_speech_duration_s is None or params.max_speech_duration_s >= 9999:
params.max_speech_duration_s = float('inf')
progress(0, desc="Filtering silent parts from audio...")
vad_options = VadOptions(
threshold=params.threshold,
min_speech_duration_ms=params.min_speech_duration_ms,
max_speech_duration_s=params.max_speech_duration_s,
min_silence_duration_ms=params.min_silence_duration_ms,
speech_pad_ms=params.speech_pad_ms
)
vad_processed, speech_chunks = self.vad.run(
audio=audio,
vad_parameters=vad_options,
progress=progress
)
try:
if vad_processed.size > 0 and speech_chunks:
if not isinstance(audio, np.ndarray):
loaded_audio = faster_whisper.decode_audio(audio, sampling_rate=self.vad.sampling_rate)
else:
loaded_audio = audio
# Convert speech_chunks to Segment objects and convert samples to seconds
segments = [Segment(start=chunk['start']/self.vad.sampling_rate, end=chunk['end']/self.vad.sampling_rate) for chunk in speech_chunks]
# merged_chunks only works on segments expressed in seconds!!
merged_chunks = merge_chunks(segments, chunk_size=300, onset=0.0, offset=None)
all_segments = []
total_elapsed_time = 0.0
for merged in merged_chunks:
chunk_start = merged['start']
chunk_end = merged['end']
# To slice audio, convert chunk_start and chunk_end from seconds to samples by mulitplying by sampling rate.
start_sample = int(chunk_start*self.vad.sampling_rate)
end_sample = int(chunk_end*self.vad.sampling_rate)
chunk_audio = loaded_audio[start_sample:end_sample]
chunk_result, chunk_time = self.transcribe(
chunk_audio,
progress,
*astuple(params)
)
# Offset timestamps
for seg in chunk_result:
seg['start'] += chunk_start
seg['end'] += chunk_start
all_segments.extend(chunk_result)
total_elapsed_time += chunk_time
result = all_segments
elapsed_time = total_elapsed_time
else:
params.vad_filter = False
except Exception as e:
print(f"Error transcribing file: {e}")
if not params.vad_filter:
result, elapsed_time = self.transcribe(
audio,
progress,
*astuple(params)
)
if bool_whisper_enable_offload:
self.offload()
if params.is_diarize:
progress(0.99, desc="Diarizing speakers...")
result, elapsed_time_diarization = self.diarizer.run(
audio=origin_audio,
use_auth_token=params.hf_token,
transcribed_result=result,
device=params.diarization_device
)
if bool_diarization_enable_offload:
self.diarizer.offload()
if not result:
print(f"Whisper did not detected any speech segments in the audio.")
result = list()
progress(1.0, desc="Processing done!")
total_elapsed_time = datetime.now() - start_time
return result, elapsed_time
def transcribe_file(self,
files_audio: Optional[List] = None,
files_video: Optional[List] = None,
files_multi: Optional[List] = None,
input_multi: str = "Audio",
input_folder_path: Optional[str] = None,
file_format: list = ["CSV"],
add_timestamp: bool = True,
translate_output: bool = False,
translate_model: str = "",
target_lang: str = "",
add_timestamp_preview: bool = False,
progress=gr.Progress(),
*whisper_params,
) -> list:
"""
Write subtitle file from Files
Parameters
----------
files_audio: list
List of files to transcribe from gr.Audio()
files_video: list
List of files to transcribe from gr.Video()
files_multi: list
List of files to transcribe from gr.Files_multi()
input_multi: str
Process single or multiple files
input_folder_path: str
Input folder path to transcribe from gr.Textbox(). If this is provided, `files` will be ignored and
this will be used instead.
file_format: str
Subtitle File format to write from gr.Dropdown(). Supported format: [CSV, SRT, TXT]
add_timestamp: bool
Boolean value from gr.Checkbox() that determines whether to add a timestamp at the end of the subtitle filename.
translate_output: bool
Translate output
translate_model: str
Translation model to use
target_lang: str
Target language to use
add_timestamp_preview: bool
Boolean value from gr.Checkbox() that determines whether to add a timestamp to output preview
progress: gr.Progress
Indicator to show progress directly in gradio.
*whisper_params: tuple
Parameters related with whisper. This will be dealt with "WhisperParameters" data class
Returns
----------
result_str:
Result of transcription to return to gr.Textbox()
result_file_path:
Output file path to return to gr.Files()
"""
try:
file_count_total = 0
process_multiple = False
files = ""
if input_multi == "Audio":
files = files_audio
elif input_multi == "Video":
files = files_video
else:
files = files_multi
file_count_total = len(files)
process_multiple = True
if input_folder_path:
files = get_media_files(input_folder_path)
if isinstance(files, str):
files = [files]
if files and isinstance(files[0], gr.utils.NamedString):
files = [file.name for file in files]
## Initialization variables & start time
files_info = {}
files_to_download = {}
time_start = datetime.now()
## Load parameters related with whisper
params = WhisperParameters.as_value(*whisper_params)
## Load model to detect language
model = whisper.load_model("base")
for file in files:
print(file)
## Detect language
mel = whisper.log_mel_spectrogram(whisper.pad_or_trim(whisper.load_audio(file))).to(model.device)
_, probs = model.detect_language(mel)
file_language = ""
file_lang_probs = ""
for key,value in whisper.tokenizer.LANGUAGES.items():
if key == str(max(probs, key=probs.get)):
file_language = value.capitalize()
for key_prob,value_prob in probs.items():
if key == key_prob:
file_lang_probs = str((round(value_prob*100)))
break
break
transcribed_segments, time_for_task = self.run(
file,
progress,
add_timestamp,
*whisper_params,
)
# Define source language
#source_lang = file_language
if params.lang == "Automatic Detection" or (params.lang).strip() == "":
source_lang = file_language
else:
source_lang = ((params.lang).strip()).capitalize()
# Translate to English using Whisper built-in functionality
transcription_note = ""
if params.is_translate:
if source_lang != "English":
transcription_note = "To English"
source_lang = "English"
else:
transcription_note = "Already in English"
# Translate the transcribed segments
translation_note = ""
if translate_output:
if source_lang != target_lang:
self.nllb_inf = NLLBInference()
if source_lang in NLLB_AVAILABLE_LANGS.keys():
transcribed_segments = self.nllb_inf.translate_text(
input_list_dict=transcribed_segments,
model_size=translate_model,
src_lang=source_lang,
tgt_lang=target_lang,
speaker_diarization=params.is_diarize
)
translation_note = "To " + target_lang
else:
translation_note = source_lang + " not supported"
else:
translation_note = "Already in " + target_lang
## Get input filename & extension
file_name, file_ext = os.path.splitext(os.path.basename(file))
## Get output as preview with or without timestamps
if add_timestamp_preview:
subtitle = get_txt(transcribed_segments)
else:
subtitle = get_plaintext(transcribed_segments)
files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task, "lang": file_language, "lang_prob": file_lang_probs, "input_source_file": (file_name+file_ext), "translation": translation_note, "transcription": transcription_note}
## Add output file as txt, srt and/or csv
for output_format in file_format:
subtitle, file_path = self.generate_and_write_file(
file_name=file_name,
transcribed_segments=transcribed_segments,
add_timestamp=add_timestamp,
file_format=output_format.lower(),
output_dir=self.output_dir
)
files_to_download[file_name+"_"+output_format.lower()] = {"path": file_path}
total_result = ""
total_info = ""
total_time = 0
file_count = 0
for file_name, info in files_info.items():
file_count += 1
# Add filename & info as first line
html_multi_tab = (" ")*3
html_single_tab = " "
html_newline = "
"
title_line = f'{html_multi_tab}Media file:{html_single_tab}{info["input_source_file"]}'
title_line += f'{html_newline}{html_multi_tab}Language:{html_single_tab}{info["lang"]} (probability {info["lang_prob"]}%)'
if params.is_translate: title_line += f'{html_newline}{html_multi_tab}Translation:{html_single_tab}{info["transcription"]} (Handled by OpenAI Whisper)'
if translate_output: title_line += f'{html_newline}{html_multi_tab}Translation:{html_single_tab}{info["translation"]} (Handled by Facebook NLLB)'
title_line = "" + title_line + "
"
if info["subtitle"].strip() == "":
info["subtitle"] = "(No speech detected)\n"
total_result += title_line+"\n"+f'{info["subtitle"]}'
total_time += info["time_for_task"]
time_end = datetime.now()
total_info = f"Processed {file_count} {'file' if file_count == 1 else 'files'} in {self.format_time((time_end-time_start).total_seconds())}"
#total_info += f"Total processing time:\t{self.format_time((time_end-time_start).total_seconds())}"
result_file_path = [info['path'] for info in files_to_download.values()]
if len(result_file_path)>0:
total_info += f" with {len(result_file_path)} output {'file' if len(result_file_path) == 1 else 'files'} available for download"
return [gr.update(value={"data": self.transform_text_to_list(total_result)},visible=True),result_file_path,total_info]
#return [result_str,result_file_path,total_info]
except Exception as e:
print(f"Error transcribing file: {e}")
finally:
self.release_cuda_memory()
def transcribe_mic(self,
mic_audio: str,
file_format: str = "SRT",
add_timestamp: bool = True,
progress=gr.Progress(),
*whisper_params,
) -> list:
"""
Write subtitle file from microphone
Parameters
----------
mic_audio: str
Audio file path from gr.Microphone()
file_format: str
Subtitle File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
add_timestamp: bool
Boolean value from gr.Checkbox() that determines whether to add a timestamp at the end of the filename.
progress: gr.Progress
Indicator to show progress directly in gradio.
*whisper_params: tuple
Parameters related with whisper. This will be dealt with "WhisperParameters" data class
Returns
----------
result_str:
Result of transcription to return to gr.Textbox()
result_file_path:
Output file path to return to gr.Files()
"""
try:
progress(0, desc="Loading Audio...")
transcribed_segments, time_for_task = self.run(
mic_audio,
progress,
add_timestamp,
*whisper_params,
)
progress(1, desc="Completed!")
subtitle, result_file_path = self.generate_and_write_file(
file_name="Mic",
transcribed_segments=transcribed_segments,
add_timestamp=add_timestamp,
file_format=file_format,
output_dir=self.output_dir
)
result_str = f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
return [result_str, result_file_path]
except Exception as e:
print(f"Error transcribing file: {e}")
finally:
self.release_cuda_memory()
def transcribe_youtube(self,
youtube_link: str,
file_format: str = "SRT",
add_timestamp: bool = True,
progress=gr.Progress(),
*whisper_params,
) -> list:
"""
Write subtitle file from Youtube
Parameters
----------
youtube_link: str
URL of the Youtube video to transcribe from gr.Textbox()
file_format: str
Subtitle File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
add_timestamp: bool
Boolean value from gr.Checkbox() that determines whether to add a timestamp at the end of the filename.
progress: gr.Progress
Indicator to show progress directly in gradio.
*whisper_params: tuple
Parameters related with whisper. This will be dealt with "WhisperParameters" data class
Returns
----------
result_str:
Result of transcription to return to gr.Textbox()
result_file_path:
Output file path to return to gr.Files()
"""
try:
progress(0, desc="Loading Audio from Youtube...")
yt = get_ytdata(youtube_link)
audio = get_ytaudio(yt)
transcribed_segments, time_for_task = self.run(
audio,
progress,
add_timestamp,
*whisper_params,
)
progress(1, desc="Completed!")
file_name = safe_filename(yt.title)
subtitle, result_file_path = self.generate_and_write_file(
file_name=file_name,
transcribed_segments=transcribed_segments,
add_timestamp=add_timestamp,
file_format=file_format,
output_dir=self.output_dir
)
result_str = f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
if os.path.exists(audio):
os.remove(audio)
return [result_str, result_file_path]
except Exception as e:
print(f"Error transcribing file: {e}")
finally:
self.release_cuda_memory()
@staticmethod
def generate_and_write_file(file_name: str,
transcribed_segments: list,
add_timestamp: bool,
file_format: str,
output_dir: str
) -> str:
"""
Writes subtitle file
Parameters
----------
file_name: str
Output file name
transcribed_segments: list
Text segments transcribed from audio
add_timestamp: bool
Determines whether to add a timestamp to the end of the filename.
file_format: str
File format to write. Supported formats: [SRT, WebVTT, txt, csv]
output_dir: str
Directory path of the output
Returns
----------
content: str
Result of the transcription
output_path: str
output file path
"""
if add_timestamp:
#timestamp = datetime.now().strftime("%m%d%H%M%S")
timestamp = datetime.now().strftime("%Y%m%d %H%M%S")
output_path = os.path.join(output_dir, f"{file_name} - {timestamp}")
else:
output_path = os.path.join(output_dir, f"{file_name}")
file_format = file_format.strip().lower()
if file_format == "srt":
content = get_srt(transcribed_segments)
output_path += '.srt'
elif file_format == "webvtt":
content = get_vtt(transcribed_segments)
output_path += '.vtt'
elif file_format == "txt":
content = get_txt(transcribed_segments)
output_path += '.txt'
elif file_format == "csv":
content = get_csv(transcribed_segments)
output_path += '.csv'
write_file(content, output_path)
return content, output_path
def offload(self):
"""Offload the model and free up the memory"""
if self.model is not None:
del self.model
self.model = None
if self.device == "cuda":
self.release_cuda_memory()
gc.collect()
@staticmethod
def transform_text_to_list(inputdata: str) -> list:
outputdata = []
temp_inputdata = (inputdata.rstrip("\n")).split("\n")
for temp_line in temp_inputdata:
temp_line_items = temp_line.split("\t")
col_time,col_speaker,col_text = "","",""
if len(temp_line_items)==1:
col_text = temp_line_items[0]
if len(temp_line_items)==2:
if (temp_line_items[0])[0].isdigit():
col_time = temp_line_items[0]
else:
col_speaker = temp_line_items[0]
col_text = temp_line_items[1]
if len(temp_line_items)==3:
col_time = temp_line_items[0]
col_speaker = temp_line_items[1]
col_text = temp_line_items[2]
outputdata.append([col_time,col_speaker,col_text])
return outputdata
@staticmethod
def format_time(elapsed_time: float) -> str:
"""
Get {hours} {minutes} {seconds} time format string
Parameters
----------
elapsed_time: str
Elapsed time for transcription
Returns
----------
Time format string
"""
hours, rem = divmod(elapsed_time, 3600)
minutes, seconds = divmod(rem, 60)
time_str = ""
hours = round(hours)
if hours:
if hours == 1:
time_str += f"{hours} hour "
else:
time_str += f"{hours} hours "
minutes = round(minutes)
if minutes:
if minutes == 1:
time_str += f"{minutes} minute "
else:
time_str += f"{minutes} minutes "
seconds = round(seconds)
if seconds == 1:
time_str += f"{seconds} second"
else:
time_str += f"{seconds} seconds"
return time_str.strip()
@staticmethod
def get_device():
if torch.cuda.is_available():
return "cuda"
elif torch.backends.mps.is_available():
if not WhisperBase.is_sparse_api_supported():
# Device `SparseMPS` is not supported for now. See : https://github.com/pytorch/pytorch/issues/87886
return "cpu"
return "mps"
else:
return "cpu"
@staticmethod
def is_sparse_api_supported():
if not torch.backends.mps.is_available():
return False
try:
device = torch.device("mps")
sparse_tensor = torch.sparse_coo_tensor(
indices=torch.tensor([[0, 1], [2, 3]]),
values=torch.tensor([1, 2]),
size=(4, 4),
device=device
)
return True
except RuntimeError:
return False
@staticmethod
def release_cuda_memory():
"""Release memory"""
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
@staticmethod
def remove_input_files(file_paths: List[str]):
"""Remove gradio cached files"""
if not file_paths:
return
for file_path in file_paths:
if file_path and os.path.exists(file_path):
os.remove(file_path)
@staticmethod
def cache_parameters(
params: WhisperValues,
file_format: str = "SRT",
add_timestamp: bool = True
):
"""Cache parameters to the yaml file"""
cached_params = load_yaml(DEFAULT_PARAMETERS_CONFIG_PATH)
param_to_cache = params.to_dict()
cached_yaml = {**cached_params, **param_to_cache}
cached_yaml["whisper"]["add_timestamp"] = add_timestamp
cached_yaml["whisper"]["file_format"] = file_format
suppress_token = cached_yaml["whisper"].get("suppress_tokens", None)
if suppress_token and isinstance(suppress_token, list):
cached_yaml["whisper"]["suppress_tokens"] = str(suppress_token)
if cached_yaml["whisper"].get("lang", None) is None:
cached_yaml["whisper"]["lang"] = AUTOMATIC_DETECTION.unwrap()
else:
language_dict = whisper.tokenizer.LANGUAGES
cached_yaml["whisper"]["lang"] = language_dict[cached_yaml["whisper"]["lang"]]
if cached_yaml["vad"].get("max_speech_duration_s", float('inf')) == float('inf'):
cached_yaml["vad"]["max_speech_duration_s"] = GRADIO_NONE_NUMBER_MAX
if cached_yaml is not None and cached_yaml:
save_yaml(cached_yaml, DEFAULT_PARAMETERS_CONFIG_PATH)
@staticmethod
def resample_audio(audio: Union[str, np.ndarray],
new_sample_rate: int = 16000,
original_sample_rate: Optional[int] = None,) -> np.ndarray:
"""Resamples audio to 16k sample rate, standard on Whisper model"""
if isinstance(audio, str):
audio, original_sample_rate = torchaudio.load(audio)
else:
if original_sample_rate is None:
raise ValueError("original_sample_rate must be provided when audio is numpy array.")
audio = torch.from_numpy(audio)
resampler = torchaudio.transforms.Resample(orig_freq=original_sample_rate, new_freq=new_sample_rate)
resampled_audio = resampler(audio).numpy()
return resampled_audio