app.py

import gradio as gr
import json
import ffmpeg
import os
from pathlib import Path
import time
from transformers import pipeline
import torch

# checkpoint = "openai/whisper-tiny"
# checkpoint = "openai/whisper-base"
checkpoint = "openai/whisper-small"

if torch.cuda.is_available() and torch.cuda.device_count() > 0:
    from transformers import (
        AutomaticSpeechRecognitionPipeline,
        WhisperForConditionalGeneration,
        WhisperProcessor,
    )
    model = WhisperForConditionalGeneration.from_pretrained(
        checkpoint).to("cuda").half()
    processor = WhisperProcessor.from_pretrained(checkpoint)
    pipe = AutomaticSpeechRecognitionPipeline(
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        batch_size=8,
        torch_dtype=torch.float16,
        device="cuda:0"
    )
else:
    pipe = pipeline(model=checkpoint)


# TODO: no longer need to set these manually once the models have been updated on the Hub
# whisper-tiny
# pipe.model.generation_config.alignment_heads = [[2, 2], [3, 0], [3, 2], [3, 3], [3, 4], [3, 5]]
# whisper-base
# pipe.model.generation_config.alignment_heads = [[3, 1], [4, 2], [4, 3], [4, 7], [5, 1], [5, 2], [5, 4], [5, 6]]
# whisper-small
pipe.model.generation_config.alignment_heads = [[5, 3], [5, 9], [
    8, 0], [8, 4], [8, 7], [8, 8], [9, 0], [9, 7], [9, 9], [10, 5]]


videos_out_path = Path("./videos_out")
videos_out_path.mkdir(parents=True, exist_ok=True)

samples_data = sorted(Path('examples').glob('*.json'))
SAMPLES = []
for file in samples_data:
    with open(file) as f:
        sample = json.load(f)
    SAMPLES.append(sample)
VIDEOS = list(map(lambda x: [x['video']], SAMPLES))


async def speech_to_text(video_in):
    """
    Takes a video path to convert to audio, transcribe audio channel to text and char timestamps

    Using https://huggingface.co/tasks/automatic-speech-recognition pipeline
    """
    video_in = video_in[0] if isinstance(video_in, list) else video_in
    if (video_in == None):
        raise ValueError("Video input undefined")

    video_path = Path(video_in.name)
    try:
        # convert video to audio 16k using PIPE to audio_memory
        audio_memory, _ = ffmpeg.input(video_path).output(
            '-', format="wav", ac=1, ar=pipe.feature_extractor.sampling_rate).overwrite_output().global_args('-loglevel', 'quiet').run(capture_stdout=True)
    except Exception as e:
        raise RuntimeError("Error converting video to audio")

    try:
        print(f'Transcribing via local model')
        output = pipe(audio_memory, chunk_length_s=10,
                      stride_length_s=[4, 2], return_timestamps="word")
        transcription = output["text"]
        chunks = output["chunks"]
        timestamps_var = [{"word": chunk["text"], "timestamp":(
            chunk["timestamp"][0], chunk["timestamp"][1]), "state": True} for chunk in chunks]

        words = [(word['word'], '+' if word['state'] else '-')
                 for word in timestamps_var]
        return (words, timestamps_var, video_in.name)
    except Exception as e:
        raise RuntimeError("Error Running inference with local model", e)


async def cut_timestamps_to_video(video_in, timestamps_var):
    video_in = video_in[0] if isinstance(video_in, list) else video_in
    if (video_in == None or timestamps_var == None):
        raise ValueError("Inputs undefined")

    video_path = Path(video_in.name)
    video_file_name = video_path.stem

    timestamps_to_cut = [
        (timestamps_var[i]['timestamp'][0], timestamps_var[i]['timestamp'][1])
        for i in range(len(timestamps_var)) if timestamps_var[i]['state']]

    between_str = '+'.join(
        map(lambda t: f'between(t,{t[0]},{t[1]})', timestamps_to_cut))

    if timestamps_to_cut:
        video_file = ffmpeg.input(video_path)
        video = video_file.video.filter(
            "select", f'({between_str})').filter("setpts", "N/FRAME_RATE/TB")
        audio = video_file.audio.filter(
            "aselect", f'({between_str})').filter("asetpts", "N/SR/TB")

        output_video = f'./videos_out/{video_file_name}.mp4'
        ffmpeg.concat(video, audio, v=1, a=1).output(
            output_video).overwrite_output().global_args('-loglevel', 'quiet').run()
    else:
        output_video = video_path

    return output_video

css = """
#words-container {
    max-height: 400px;
    overflow-y: scroll !important;
}
"""
with gr.Blocks(css=css) as demo:
    timestamps_var = gr.JSON(visible=False)
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            # Whisper: Word-Level Video Trimming
            Quick edit a video by trimming out words.
            Using the [Huggingface Automatic Speech Recognition Pipeline](https://huggingface.co/tasks/automatic-speech-recognition)
            with [Whisper](https://huggingface.co/docs/transformers/model_doc/whisper)
            """)

    with gr.Row():
        with gr.Column():
            file_upload = gr.File(
                label="Upload Video File", file_count=1, scale=1)
            video_preview = gr.Video(
                label="Video Preview", scale=3, interactive=False)
            # with gr.Row():
            #     transcribe_btn = gr.Button(
            #         "Transcribe Audio")

        with gr.Column():
            text_in = gr.HighlightedText(
                label="Transcription", combine_adjacent=False, show_legend=True, color_map={"+": "green", "-": "red"}, elem_id="words-container")
            with gr.Row():
                cut_btn = gr.Button("Cut Video")
                select_all_words = gr.Button("Select All Words")
                reset_words = gr.Button("Reset Words")
            video_out = gr.Video(label="Video Out")
    with gr.Row():
        gr.Examples(
            fn=speech_to_text,
            examples=["./examples/ShiaLaBeouf.mp4",
                      "./examples/zuckyuval.mp4",
                      "./examples/cooking.mp4"],
            inputs=[file_upload],
            outputs=[text_in, timestamps_var, video_preview],
            cache_examples=True)

    with gr.Row():
        gr.Markdown("""
        #### Video Credits

        1. [Cooking](https://vimeo.com/573792389)
        1. [Shia LaBeouf "Just Do It"](https://www.youtube.com/watch?v=n2lTxIk_Dr0)
        1. [Mark Zuckerberg & Yuval Noah Harari in Conversation](https://www.youtube.com/watch?v=Boj9eD0Wug8)
        """)

    def select_text(evt: gr.SelectData, timestamps_var):
        index = evt.index
        timestamps_var[index]['state'] = not timestamps_var[index]['state']
        words = [(word['word'], '+' if word['state'] else '-')
                 for word in timestamps_var]
        return timestamps_var, words

    def words_selection(timestamps_var, reset=False):
        if reset:
            for word in timestamps_var:
                word['state'] = True
        else:
            # reverse the state of all words
            for word in timestamps_var:
                word['state'] = False

        words = [(word['word'], '+' if word['state'] else '-')
                 for word in timestamps_var]
        return timestamps_var, words

    file_upload.upload(speech_to_text, inputs=[file_upload], outputs=[
        text_in, timestamps_var, video_preview])
    select_all_words.click(words_selection, inputs=[timestamps_var], outputs=[
                           timestamps_var, text_in], queue=False, show_progress=False)
    reset_words.click(lambda x: words_selection(x, True), inputs=[timestamps_var], outputs=[
                      timestamps_var, text_in], queue=False, show_progress=False)
    text_in.select(select_text, inputs=timestamps_var,
                   outputs=[timestamps_var, text_in], queue=False, show_progress=False)
    # transcribe_btn.click(speech_to_text, inputs=[file_upload], outputs=[
    #                      text_in, transcription_var, timestamps_var, video_preview])
    cut_btn.click(cut_timestamps_to_video, [
                  file_upload, timestamps_var], [video_out])

demo.queue()
if __name__ == "__main__":
    demo.launch(debug=True)