Spaces:
Runtime error
Runtime error
| import logging | |
| import warnings | |
| import gradio as gr | |
| import pytube as pt | |
| import torch | |
| from huggingface_hub import model_info | |
| from transformers import pipeline | |
| from transformers.utils.logging import disable_progress_bar | |
| warnings.filterwarnings("ignore") | |
| disable_progress_bar() | |
| DEFAULT_MODEL_NAME = "bofenghuang/whisper-large-v2-cv11-german" | |
| # make sure no OOM | |
| MODEL_NAMES = [ | |
| "bofenghuang/whisper-medium-cv11-german", | |
| "bofenghuang/whisper-large-v2-cv11-german", | |
| ] | |
| LANG = "de" | |
| CHUNK_LENGTH_S = 30 | |
| MAX_NEW_TOKENS = 225 | |
| logging.basicConfig( | |
| format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s", | |
| datefmt="%Y-%m-%dT%H:%M:%SZ", | |
| ) | |
| logger = logging.getLogger(__name__) | |
| logger.setLevel(logging.DEBUG) | |
| device = 0 if torch.cuda.is_available() else "cpu" | |
| logger.info(f"Model will be loaded on device `{device}`") | |
| cached_models = {} | |
| def print_cuda_memory_info(): | |
| used_mem, tot_mem = torch.cuda.mem_get_info() | |
| logger.info(f"CUDA memory info - Free: {used_mem / 1024 ** 3:.2f} Gb, used: {(tot_mem - used_mem) / 1024 ** 3:.2f} Gb, total: {tot_mem / 1024 ** 3:.2f} Gb") | |
| def print_memory_info(): | |
| # todo | |
| if device == "cpu": | |
| pass | |
| else: | |
| print_cuda_memory_info() | |
| def maybe_load_cached_pipeline(model_name): | |
| pipe = cached_models.get(model_name) | |
| if pipe is None: | |
| # load pipeline | |
| # todo: set decoding option for pipeline | |
| pipe = pipeline( | |
| task="automatic-speech-recognition", | |
| model=model_name, | |
| chunk_length_s=CHUNK_LENGTH_S, | |
| device=device, | |
| ) | |
| # set forced_decoder_ids | |
| pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=LANG, task="transcribe") | |
| # limit genneration max length | |
| pipe.model.config.max_length = MAX_NEW_TOKENS + 1 | |
| logger.info(f"`{model_name}` pipeline has been initialized") | |
| print_memory_info() | |
| cached_models[model_name] = pipe | |
| return pipe | |
| def transcribe(microphone, file_upload, model_name): | |
| warn_output = "" | |
| if (microphone is not None) and (file_upload is not None): | |
| warn_output = ( | |
| "WARNING: You've uploaded an audio file and used the microphone. " | |
| "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n" | |
| ) | |
| elif (microphone is None) and (file_upload is None): | |
| return "ERROR: You have to either use the microphone or upload an audio file" | |
| file = microphone if microphone is not None else file_upload | |
| pipe = maybe_load_cached_pipeline(model_name) | |
| text = pipe(file)["text"] | |
| logger.info(f"Transcription by `{model_name}`: {text}") | |
| return warn_output + text | |
| def _return_yt_html_embed(yt_url): | |
| video_id = yt_url.split("?v=")[-1] | |
| HTML_str = ( | |
| f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>' | |
| " </center>" | |
| ) | |
| return HTML_str | |
| def yt_transcribe(yt_url, model_name): | |
| yt = pt.YouTube(yt_url) | |
| html_embed_str = _return_yt_html_embed(yt_url) | |
| stream = yt.streams.filter(only_audio=True)[0] | |
| stream.download(filename="audio.mp3") | |
| pipe = maybe_load_cached_pipeline(model_name) | |
| text = pipe("audio.mp3")["text"] | |
| logger.info(f"Transcription: {text}") | |
| return html_embed_str, text | |
| # load default model | |
| maybe_load_cached_pipeline(DEFAULT_MODEL_NAME) | |
| demo = gr.Blocks() | |
| mf_transcribe = gr.Interface( | |
| fn=transcribe, | |
| inputs=[ | |
| gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Record"), | |
| gr.inputs.Audio(source="upload", type="filepath", optional=True, label="Upload File"), | |
| gr.inputs.Dropdown(choices=MODEL_NAMES, default=DEFAULT_MODEL_NAME, label="Whisper Model"), | |
| ], | |
| # outputs="text", | |
| outputs=gr.outputs.Textbox(label="Transcription"), | |
| layout="horizontal", | |
| theme="huggingface", | |
| title="Whisper German Demo 🇩🇪 : Transcribe Audio", | |
| description="Transcribe long-form microphone or audio inputs with the click of a button!", | |
| allow_flagging="never", | |
| ) | |
| yt_transcribe = gr.Interface( | |
| fn=yt_transcribe, | |
| inputs=[ | |
| gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"), | |
| gr.inputs.Dropdown(choices=MODEL_NAMES, default=DEFAULT_MODEL_NAME, label="Whisper Model"), | |
| ], | |
| # outputs=["html", "text"], | |
| outputs=[ | |
| gr.outputs.HTML(label="YouTube Page"), | |
| gr.outputs.Textbox(label="Transcription"), | |
| ], | |
| layout="horizontal", | |
| theme="huggingface", | |
| title="Whisper German Demo 🇩🇪 : Transcribe YouTube", | |
| description="Transcribe long-form YouTube videos with the click of a button!", | |
| allow_flagging="never", | |
| ) | |
| with demo: | |
| gr.TabbedInterface([mf_transcribe, yt_transcribe], ["Transcribe Audio", "Transcribe YouTube"]) | |
| # demo.launch(server_name="0.0.0.0", debug=True, share=True) | |
| demo.launch(enable_queue=True) | |