run_demo_openai_layout.py

#! /usr/bin/env python
# coding=utf-8
# Copyright 2022 Bofeng Huang

import datetime
import logging
import os
import re
import warnings

import gradio as gr
import pandas as pd
import psutil
import pytube as pt
import torch
import whisper
from huggingface_hub import hf_hub_download, model_info
from transformers.utils.logging import disable_progress_bar

import nltk
nltk.download("punkt")

from nltk.tokenize import sent_tokenize

warnings.filterwarnings("ignore")
disable_progress_bar()

DEFAULT_MODEL_NAME = "bofenghuang/whisper-large-v2-cv11-german"
CHECKPOINT_FILENAME = "checkpoint_openai.pt"

GEN_KWARGS = {
    "task": "transcribe",
    "language": "de",
    # "without_timestamps": True,
    # decode options
    # "beam_size": 5,
    # "patience": 2,
    # disable fallback
    # "compression_ratio_threshold": None,
    # "logprob_threshold": None,
    # vad threshold
    # "no_speech_threshold": None,
}

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"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logger.info(f"Model will be loaded on device `{device}`")

cached_models = {}


def format_timestamp(seconds):
    return str(datetime.timedelta(seconds=round(seconds)))


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 download_audio_from_youtube(yt_url, downloaded_filename="audio.wav"):
    yt = pt.YouTube(yt_url)
    stream = yt.streams.filter(only_audio=True)[0]
    # stream.download(filename="audio.mp3")
    stream.download(filename=downloaded_filename)
    return downloaded_filename


def download_video_from_youtube(yt_url, downloaded_filename="video.mp4"):
    yt = pt.YouTube(yt_url)
    stream = yt.streams.filter(progressive=True, file_extension="mp4").order_by("resolution").desc().first()
    stream.download(filename=downloaded_filename)
    logger.info(f"Download YouTube video from {yt_url}")
    return downloaded_filename


def _print_memory_info():
    memory = psutil.virtual_memory()
    logger.info(
        f"Memory info - Free: {memory.available / (1024 ** 3):.2f} Gb, used: {memory.percent}%, total: {memory.total / (1024 ** 3):.2f} Gb"
    )


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():
    _print_memory_info()
    _print_cuda_memory_info()


def maybe_load_cached_pipeline(model_name):
    model = cached_models.get(model_name)
    if model is None:
        downloaded_model_path = hf_hub_download(repo_id=model_name, filename=CHECKPOINT_FILENAME)

        model = whisper.load_model(downloaded_model_path, device=device)
        logger.info(f"`{model_name}` has been loaded on device `{device}`")

        print_memory_info()

        cached_models[model_name] = model
    return model


def infer(model, filename, with_timestamps, return_df=False):
    if with_timestamps:
        model_outputs = model.transcribe(filename, **GEN_KWARGS)
        if return_df:
            model_outputs_df = pd.DataFrame(model_outputs["segments"])
            # print(model_outputs)
            # print(model_outputs_df)
            # print(model_outputs_df.info(verbose=True))
            model_outputs_df = model_outputs_df[["start", "end", "text"]]
            model_outputs_df["start"] = model_outputs_df["start"].map(format_timestamp)
            model_outputs_df["end"] = model_outputs_df["end"].map(format_timestamp)
            model_outputs_df["text"] = model_outputs_df["text"].str.strip()
            return model_outputs_df
        else:
            return "\n\n".join(
                [
                    f'Segment {segment["id"]+1} from {segment["start"]:.2f}s to {segment["end"]:.2f}s:\n{segment["text"].strip()}'
                    for segment in model_outputs["segments"]
                ]
            )
    else:
        text = model.transcribe(filename, without_timestamps=True, **GEN_KWARGS)["text"]
        if return_df:
            return pd.DataFrame({"text": sent_tokenize(text)})
        else:
            return text


def transcribe(microphone, file_upload, with_timestamps, model_name=DEFAULT_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

    model = maybe_load_cached_pipeline(model_name)
    # text = model.transcribe(file, **GEN_KWARGS)["text"]
    # text = infer(model, file, with_timestamps)
    text = infer(model, file, with_timestamps, return_df=True)

    logger.info(f'Transcription by `{model_name}`:\n{text.to_json(orient="index", force_ascii=False, indent=2)}\n')

    # return warn_output + text
    return text


def yt_transcribe(yt_url, with_timestamps, model_name=DEFAULT_MODEL_NAME):
    # html_embed_str = _return_yt_html_embed(yt_url)
    audio_file_path = download_audio_from_youtube(yt_url)

    model = maybe_load_cached_pipeline(model_name)
    # text = model.transcribe("audio.mp3", **GEN_KWARGS)["text"]
    # text = infer(model, audio_file_path, with_timestamps)
    text = infer(model, audio_file_path, with_timestamps, return_df=True)

    logger.info(f'Transcription by `{model_name}` of "{yt_url}":\n{text.to_json(orient="index", force_ascii=False, indent=2)}\n')

    # return html_embed_str, text
    return text


def video_transcribe(video_file_path, with_timestamps, model_name=DEFAULT_MODEL_NAME):
    if video_file_path is None:
        raise ValueError("Failed to transcribe video as no video_file_path has been defined")

    audio_file_path = re.sub(r"\.mp4$", ".wav", video_file_path)
    os.system(f'ffmpeg -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{audio_file_path}"')

    model = maybe_load_cached_pipeline(model_name)
    # text = model.transcribe("audio.mp3", **GEN_KWARGS)["text"]
    text = infer(model, audio_file_path, with_timestamps, return_df=True)

    logger.info(f'Transcription by `{model_name}`:\n{text.to_json(orient="index", force_ascii=False, indent=2)}\n')

    return text


# load default model
maybe_load_cached_pipeline(DEFAULT_MODEL_NAME)

# default_text_output_df = pd.DataFrame(columns=["start", "end", "text"])
default_text_output_df = pd.DataFrame(columns=["text"])

with gr.Blocks() as demo:

    with gr.Tab("Transcribe Audio"):
        gr.Markdown(
            f"""
              <div>
              <h1 style='text-align: center'>Whisper German Demo 🇩🇪 : Transcribe Audio</h1>
              </div>
              Transcribe long-form microphone or audio inputs!

              Demo uses the fine-tuned checkpoint: <a href='https://huggingface.co/{DEFAULT_MODEL_NAME}' target='_blank'><b>{DEFAULT_MODEL_NAME}</b></a> to transcribe audio files of arbitrary length.
          """
        )

        microphone_input = gr.inputs.Audio(source="microphone", type="filepath", label="Record", optional=True)
        upload_input = gr.inputs.Audio(source="upload", type="filepath", label="Upload File", optional=True)
        with_timestamps_input = gr.Checkbox(label="With timestamps?")

        microphone_transcribe_btn = gr.Button("Transcribe Audio")

        # gr.Markdown('''
        #     Here you will get generated transcrit.
        # ''')

        # microphone_text_output = gr.outputs.Textbox(label="Transcription")
        text_output_df2 = gr.DataFrame(
            value=default_text_output_df,
            label="Transcription",
            row_count=(0, "dynamic"),
            max_rows=10,
            wrap=True,
            overflow_row_behaviour="paginate",
        )

        microphone_transcribe_btn.click(
            transcribe, inputs=[microphone_input, upload_input, with_timestamps_input], outputs=text_output_df2
        )

    # with gr.Tab("Transcribe YouTube"):
    #     gr.Markdown(
    #         f"""
    #           <div>
    #           <h1 style='text-align: center'>Whisper German Demo 🇩🇪 : Transcribe YouTube</h1>
    #           </div>
    #           Transcribe long-form YouTube videos!

    #           Demo uses the fine-tuned checkpoint: <a href='https://huggingface.co/{DEFAULT_MODEL_NAME}' target='_blank'><b>{DEFAULT_MODEL_NAME}</b></a> to transcribe video files of arbitrary length.
    #       """
    #     )

    #     yt_link_input2 = gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL")
    #     with_timestamps_input2 = gr.Checkbox(label="With timestamps?", value=True)

    #     yt_transcribe_btn = gr.Button("Transcribe YouTube")

    #     # yt_text_output = gr.outputs.Textbox(label="Transcription")
    #     text_output_df3 = gr.DataFrame(
    #         value=default_text_output_df,
    #         label="Transcription",
    #         row_count=(0, "dynamic"),
    #         max_rows=10,
    #         wrap=True,
    #         overflow_row_behaviour="paginate",
    #     )
    #     # yt_html_output = gr.outputs.HTML(label="YouTube Page")

    #     yt_transcribe_btn.click(yt_transcribe, inputs=[yt_link_input2, with_timestamps_input2], outputs=[text_output_df3])

    with gr.Tab("Transcribe Video"):
        gr.Markdown(
            f"""
              <div>
              <h1 style='text-align: center'>Whisper German Demo 🇩🇪 : Transcribe Video</h1>
              </div>
              Transcribe long-form YouTube videos or uploaded video inputs!

              Demo uses the fine-tuned checkpoint: <a href='https://huggingface.co/{DEFAULT_MODEL_NAME}' target='_blank'><b>{DEFAULT_MODEL_NAME}</b></a> to transcribe video files of arbitrary length.
          """
        )

        yt_link_input = gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL")
        download_youtube_btn = gr.Button("Download Youtube video")
        downloaded_video_output = gr.Video(label="Video file", mirror_webcam=False)
        download_youtube_btn.click(download_video_from_youtube, inputs=[yt_link_input], outputs=[downloaded_video_output])

        with_timestamps_input3 = gr.Checkbox(label="With timestamps?", value=True)
        video_transcribe_btn = gr.Button("Transcribe video")
        text_output_df = gr.DataFrame(
            value=default_text_output_df,
            label="Transcription",
            row_count=(0, "dynamic"),
            max_rows=10,
            wrap=True,
            overflow_row_behaviour="paginate",
        )

        video_transcribe_btn.click(video_transcribe, inputs=[downloaded_video_output, with_timestamps_input3], outputs=[text_output_df])

# demo.launch(server_name="0.0.0.0", debug=True)
# demo.launch(server_name="0.0.0.0", debug=True, share=True)
demo.launch(enable_queue=True)