jhj0517
Initial commit
c62f556
import whisper
from modules.subtitle_manager import get_srt,get_vtt,safe_filename
from modules.youtube_manager import get_ytdata,get_ytaudio
import gradio as gr
import os
DEFAULT_MODEL_SIZE="tiny"
class WhisperInference():
def __init__(self):
print("\nInitializing Model..\n")
self.current_model_size = DEFAULT_MODEL_SIZE
self.model = whisper.load_model(name=DEFAULT_MODEL_SIZE,download_root="models")
self.available_models = ["tiny","tiny.en"]
self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
def transcribe_file(self,fileobjs
,model_size,lang,subformat,istranslate,
progress=gr.Progress()):
def progress_callback(progress_value):
progress(progress_value,desc="Transcribing..")
if model_size != self.current_model_size:
progress(0,desc="Initializing Model..")
self.current_model_size = model_size
self.model = whisper.load_model(name=model_size,download_root="models")
if lang == "Automatic Detection" :
lang = None
progress(0,desc="Loading Audio..")
files_info = {}
for fileobj in fileobjs:
print(f"\n\n {fileobj.name} \n\n")
audio = whisper.load_audio(fileobj.name)
translatable_model = ["large","large-v1","large-v2"]
if istranslate and self.current_model_size in translatable_model:
result = self.model.transcribe(audio=audio,language=lang,verbose=False,task="translate",progress_callback=progress_callback)
else :
result = self.model.transcribe(audio=audio,language=lang,verbose=False,progress_callback=progress_callback)
progress(1,desc="Completed!")
file_name, file_ext = os.path.splitext(os.path.basename(fileobj.orig_name))
file_name = file_name[:-9]
file_name = safe_filename(file_name)
if subformat == "SRT":
subtitle = get_srt(result["segments"])
elif subformat == "WebVTT":
subtitle = get_vtt(result["segments"])
files_info[file_name] = subtitle
total_result = ''
for file_name,subtitle in files_info.items():
total_result+='------------------------------------\n'
total_result+=f'{file_name}\n\n'
total_result+=f'{subtitle}'
return f"\n\n{total_result}"
def transcribe_youtube(self,youtubelink
,model_size,lang,subformat,istranslate,
progress=gr.Progress()):
def progress_callback(progress_value):
progress(progress_value,desc="Transcribing..")
if model_size != self.current_model_size:
progress(0,desc="Initializing Model..")
self.current_model_size = model_size
self.model = whisper.load_model(name=model_size,download_root="models")
if lang == "Automatic Detection" :
lang = None
progress(0,desc="Loading Audio from Youtube..")
yt = get_ytdata(youtubelink)
audio = whisper.load_audio(get_ytaudio(yt))
translatable_model = ["large","large-v1","large-v2"]
if istranslate and self.current_model_size in translatable_model:
result = self.model.transcribe(audio=audio,language=lang,verbose=False,task="translate",progress_callback=progress_callback)
else :
result = self.model.transcribe(audio=audio,language=lang,verbose=False,progress_callback=progress_callback)
progress(1,desc="Completed!")
file_name = safe_filename(yt.title)
if subformat == "SRT":
subtitle = get_srt(result["segments"])
elif subformat == "WebVTT":
subtitle = get_vtt(result["segments"])
return f"\n\n{subtitle}"
def transcribe_mic(self,micaudio
,model_size,lang,subformat,istranslate,
progress=gr.Progress()):
def progress_callback(progress_value):
progress(progress_value,desc="Transcribing..")
if model_size != self.current_model_size:
progress(0,desc="Initializing Model..")
self.current_model_size = model_size
self.model = whisper.load_model(name=model_size,download_root="models")
if lang == "Automatic Detection" :
lang = None
progress(0,desc="Loading Audio..")
translatable_model = ["large","large-v1","large-v2"]
if istranslate and self.current_model_size in translatable_model:
result = self.model.transcribe(audio=micaudio,language=lang,verbose=False,task="translate",progress_callback=progress_callback)
else :
result = self.model.transcribe(audio=micaudio,language=lang,verbose=False,progress_callback=progress_callback)
progress(1,desc="Completed!")
if subformat == "SRT":
subtitle = get_srt(result["segments"])
elif subformat == "WebVTT":
subtitle = get_vtt(result["segments"])
return f"\n\n{subtitle}"