Create app.py

app.py

@@ -0,0 +1,329 @@
+import gradio as gr
+import json
+from difflib import Differ
+import ffmpeg
+import os
+from pathlib import Path
+import time
+import aiohttp
+import asyncio
+from transformers import pipeline
+
+# Set true if you're using huggingface inference API API https://huggingface.co/inference-api
+API_BACKEND = True
+MODEL = "facebook/wav2vec2-base-960h"
+
+if API_BACKEND:
+    from dotenv import load_dotenv
+    import base64
+    import asyncio
+    load_dotenv(Path(".env"))
+
+    HF_TOKEN = os.environ["HF_TOKEN"] = ""
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    API_URL = f'https://api-inference.huggingface.co/models/{MODEL}'
+
+else:
+    import torch
+    from transformers import pipeline
+
+    # is cuda available?
+    cuda = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    device = 0 if torch.cuda.is_available() else -1
+    speech_recognizer = pipeline(
+        task="automatic-speech-recognition",
+        model=MODEL,
+        tokenizer=MODEL,
+        framework="pt",
+        device=device,
+    )
+
+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))
+
+total_inferences_since_reboot = 0
+total_cuts_since_reboot = 0
+
+
+async def speech_to_text(video_file_path):
+    """
+    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
+    """
+    global total_inferences_since_reboot
+    if (video_file_path == None):
+        raise ValueError("Error no video input")
+
+    video_path = Path(video_file_path)
+    try:
+        # convert video to audio 16k using PIPE to audio_memory
+        audio_memory, _ = ffmpeg.input(video_path).output(
+            '-', format="wav", ac=1, ar='16k').overwrite_output().global_args('-loglevel', 'quiet').run(capture_stdout=True)
+    except Exception as e:
+        raise RuntimeError("Error converting video to audio")
+
+    ping("speech_to_text")
+    last_time = time.time()
+    if API_BACKEND:
+        # Using Inference API https://huggingface.co/inference-api
+        # try twice, because the model must be loaded
+        for i in range(10):
+            for tries in range(4):
+                print(f'Transcribing from API attempt {tries}')
+                try:
+                    inference_reponse = await query_api(audio_memory)
+                    print(inference_reponse)
+                    transcription = inference_reponse["text"].lower()
+                    timestamps = [[chunk["text"].lower(), chunk["timestamp"][0], chunk["timestamp"][1]]
+                                  for chunk in inference_reponse['chunks']]
+
+                    total_inferences_since_reboot += 1
+                    print("\n\ntotal_inferences_since_reboot: ",
+                          total_inferences_since_reboot, "\n\n")
+                    return (transcription, transcription, timestamps)
+                except Exception as e:
+                    print(e)
+                    if 'error' in inference_reponse and 'estimated_time' in inference_reponse:
+                        wait_time = inference_reponse['estimated_time']
+                        print("Waiting for model to load....", wait_time)
+                        # wait for loading model
+                        # 5 seconds plus for certainty
+                        await asyncio.sleep(wait_time + 5.0)
+                    elif 'error' in inference_reponse:
+                        raise RuntimeError("Error Fetching API",
+                                           inference_reponse['error'])
+                    else:
+                        break
+            else:
+                raise RuntimeError(inference_reponse, "Error Fetching API")
+    else:
+        try:
+            print(f'Transcribing via local model')
+            output = speech_recognizer(
+                audio_memory, return_timestamps="char", chunk_length_s=10, stride_length_s=(4, 2))
+
+            transcription = output["text"].lower()
+            timestamps = [[chunk["text"].lower(), chunk["timestamp"][0].tolist(), chunk["timestamp"][1].tolist()]
+                          for chunk in output['chunks']]
+            total_inferences_since_reboot += 1
+
+            print("\n\ntotal_inferences_since_reboot: ",
+                  total_inferences_since_reboot, "\n\n")
+            return (transcription, transcription, timestamps)
+        except Exception as e:
+            raise RuntimeError("Error Running inference with local model", e)
+
+
+async def cut_timestamps_to_video(video_in, transcription, text_in, timestamps, words_to_cut):
+    """
+    Given original video input, text transcript + timestamps,
+    and edit text, cuts video segments into a single video
+    """
+    global total_cuts_since_reboot
+
+    video_path = Path(video_in)
+    video_file_name = video_path.stem
+    if (video_in == None or text_in == None or transcription == None):
+        raise ValueError("Inputs undefined")
+
+    d = Differ()
+    # compare original transcription with edit text
+    diff_chars = d.compare(transcription, text_in)
+
+    # Include additions in the filtered list
+    filtered = list(filter(lambda x: x[0] != '-' and x[0] != '+', diff_chars))
+
+    # Update grouping logic to handle additions and word cuts
+    idx = 0
+    grouped = {}
+    word_cuts = []
+    for (a, b) in zip(filtered, timestamps):
+        if a[0] != '-':
+            if idx in grouped:
+                grouped[idx].append(b)
+            else:
+                grouped[idx] = []
+                grouped[idx].append(b)
+        elif a[0] == '-':
+            idx += 1
+        elif a[0] == '+':
+            word_cuts.append(b)
+
+    # after grouping, gets the lower and upper start and time for each group
+    timestamps_to_cut = [[v[0][1], v[-1][2]] for v in grouped.values()]
+
+    # Add word cut timestamps
+    timestamps_to_cut.extend(word_cuts)
+
+    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_in)
+        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_in
+
+    tokens = [(token[2:], token[0] if token[0] != " " else None)
+              for token in filtered]
+
+    total_cuts_since_reboot += 1
+    ping("video_cuts")
+    print("\n\ntotal_cuts_since_reboot: ", total_cuts_since_reboot, "\n\n")
+    return (tokens, output_video)
+
+
+async def query_api(audio_bytes: bytes):
+    """
+    Query for Huggingface Inference API for Automatic Speech Recognition task
+    """
+    payload = json.dumps({
+        "inputs": base64.b64encode(audio_bytes).decode("utf-8"),
+        "parameters": {
+            "return_timestamps": "char",
+            "chunk_length_s": 10,
+            "stride_length_s": [4, 2]
+        },
+        "options": {"use_gpu": False}
+    }).encode("utf-8")
+    async with aiohttp.ClientSession() as session:
+        async with session.post(API_URL, headers=headers, data=payload) as response:
+            print("API Response: ", response.status)
+            if response.headers['Content-Type'] == 'application/json':
+                return await response.json()
+            elif response.headers['Content-Type'] == 'application/octet-stream':
+                return await response.read()
+            elif response.headers['Content-Type'] == 'text/plain':
+                return await response.text()
+            else:
+                raise RuntimeError("Error Fetching API")
+
+
+def ping(name):
+    url = f'https://huggingface.co/api/telemetry/spaces/radames/edit-video-by-editing-text/{name}'
+    print("ping: ", url)
+
+    async def req():
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                print("pong: ", response.status)
+    asyncio.create_task(req())
+
+
+def find_words_to_cut(transcription, words):
+    """
+    Find timestamps of words to cut in the transcription
+    """
+    timestamps_to_cut = []
+    for word in words:
+        word_lower = word.lower()
+        word_positions = [i for i, x in enumerate(transcription.split()) if x.lower() == word_lower]
+        for pos in word_positions:
+            timestamps_to_cut.append([transcription[pos:].find(' '), pos + 1])
+
+    return timestamps_to_cut
+
+
+# ---- Gradio Layout -----
+video_in = gr.Video(label="Video file", elem_id="video-container")
+text_in = gr.Textbox(label="Transcription", lines=10, interactive=True)
+words_to_cut_input = gr.Textbox(label="Words to Cut (comma-separated)", lines=1, default="")
+video_out = gr.Video(label="Video Out")
+diff_out = gr.HighlightedText(label="Cuts Diffs", combine_adjacent=True)
+examples = gr.Dataset(components=[video_in], samples=VIDEOS, type="index")
+
+css = """
+#cut_btn, #reset_btn { align-self:stretch; }
+#\\31 3 { max-width: 540px; }
+.output-markdown {max-width: 65ch !important;}
+#video-container{
+    max-width: 40rem;
+}
+"""
+with gr.Blocks(css=css) as demo:
+    transcription_var = gr.State()
+    timestamps_var = gr.State()
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("""
+            # Edit Video By Editing Text
+            This project is a quick proof of concept of a simple video editor where the edits
+            are made by editing the audio transcription.
+            Using the [Huggingface Automatic Speech Recognition Pipeline](https://huggingface.co/tasks/automatic-speech-recognition)
+            with a fine-tuned model, you can predict the text transcription and character or word-based timestamps.
+            """)
+
+    with gr.Row():
+        examples.render()
+
+        def load_example(id):
+            video = SAMPLES[id]['video']
+            transcription = SAMPLES[id]['transcription'].lower()
+            timestamps = SAMPLES[id]['timestamps']
+
+            return (video, transcription, transcription, timestamps)
+
+        examples.click(
+            load_example,
+            inputs=[examples],
+            outputs=[video_in, text_in, transcription_var, timestamps_var],
+            queue=False)
+
+    with gr.Row():
+        with gr.Column():
+            video_in.render()
+            transcribe_btn = gr.Button("Transcribe Audio")
+            transcribe_btn.click(speech_to_text, [video_in], [
+                text_in, transcription_var, timestamps_var])
+
+    with gr.Row():
+        gr.Markdown("""
+        ### Now edit as text
+        After running the video transcription, you can make cuts to the text below (only cuts, not additions!)
+        You can also specify words to cut in the "Words to Cut" input box.
+        """)
+
+    with gr.Row():
+        with gr.Column():
+            text_in.render()
+            words_to_cut_input.render()
+            with gr.Row():
+                cut_btn = gr.Button("Cut to video", elem_id="cut_btn")
+                # send audio path and hidden variables
+                cut_btn.click(cut_timestamps_to_video, [
+                    video_in, transcription_var, text_in, timestamps_var, words_to_cut_input], [diff_out, video_out])
+
+                reset_transcription = gr.Button(
+                    "Reset to last transcription", elem_id="reset_btn")
+                reset_transcription.click(
+                    lambda x: x, transcription_var, text_in)
+        with gr.Column():
+            video_out.render()
+            diff_out.render()
+
+    with gr.Row():
+        gr.Markdown("""
+        #### Video Credits
+
+        1. [Cooking](https://vimeo.com/573792389)
+        2. [Shia LaBeouf "Just Do It"](https://www.youtube.com/watch?v=n2lTxIk_Dr0)
+        3. [Mark Zuckerberg & Yuval Noah Harari in Conversation](https://www.youtube.com/watch?v=Boj9eD0Wug8)
+        """)
+
+demo.launch(debug=True)