Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import librosa | |
| import numpy as np | |
| import moviepy.editor as mpy | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import pipeline | |
| fps = 25 | |
| max_duration = 60 # seconds | |
| video_width = 640 | |
| video_height = 480 | |
| margin_left = 20 | |
| margin_right = 20 | |
| margin_top = 20 | |
| line_height = 44 | |
| background_image = Image.open("background.png") | |
| font = ImageFont.truetype("Lato-Regular.ttf", 40) | |
| text_color = (255, 200, 200) | |
| highlight_color = (255, 255, 255) | |
| # checkpoint = "openai/whisper-tiny" | |
| # checkpoint = "openai/whisper-base" | |
| checkpoint = "openai/whisper-small" | |
| pipe = pipeline(model=checkpoint) | |
| # TODO: no longer need to set these manually once the models have been updated on the Hub | |
| # whisper-base | |
| # pipe.model.config.alignment_heads = [[3, 1], [4, 2], [4, 3], [4, 7], [5, 1], [5, 2], [5, 4], [5, 6]] | |
| # whisper-small | |
| pipe.model.config.alignment_heads = [[5, 3], [5, 9], [8, 0], [8, 4], [8, 7], [8, 8], [9, 0], [9, 7], [9, 9], [10, 5]] | |
| chunks = [] | |
| def make_frame(t): | |
| global chunks | |
| # TODO speed optimization: could cache the last image returned and if the | |
| # active chunk and active word didn't change, use that last image instead | |
| # of drawing the exact same thing again | |
| # TODO in the Henry V example, the word "desires" has an ending timestamp | |
| # that's too far into the future, and so the word stays highlighted. | |
| # Could fix this by finding the latest word that is active in the chunk | |
| # and only highlight that one. | |
| image = background_image.copy() | |
| draw = ImageDraw.Draw(image) | |
| # for debugging: draw frame time | |
| #draw.text((20, 20), str(t), fill=text_color, font=font) | |
| space_length = draw.textlength(" ", font) | |
| x = margin_left | |
| y = margin_top | |
| for chunk in chunks: | |
| chunk_start = chunk["timestamp"][0] | |
| chunk_end = chunk["timestamp"][1] | |
| if chunk_end is None: chunk_end = max_duration | |
| if chunk_start <= t <= chunk_end: | |
| words = [x["text"] for x in chunk["words"]] | |
| word_times = [x["timestamp"] for x in chunk["words"]] | |
| for (word, times) in zip(words, word_times): | |
| word_length = draw.textlength(word + " ", font) - space_length | |
| if x + word_length >= video_width - margin_right: | |
| x = margin_left | |
| y += line_height | |
| if times[0] <= t <= times[1]: | |
| color = highlight_color | |
| draw.rectangle([x, y + line_height, x + word_length, y + line_height + 4], fill=color) | |
| else: | |
| color = text_color | |
| draw.text((x, y), word, fill=color, font=font) | |
| x += word_length + space_length | |
| break | |
| return np.array(image) | |
| def predict(audio_path): | |
| global chunks | |
| audio_data, sr = librosa.load(audio_path, mono=True) | |
| duration = librosa.get_duration(y=audio_data, sr=sr) | |
| duration = min(max_duration, duration) | |
| audio_data = audio_data[:int(duration * sr)] | |
| # Run Whisper to get word-level timestamps. | |
| audio_inputs = librosa.resample(audio_data, orig_sr=sr, target_sr=pipe.feature_extractor.sampling_rate) | |
| output = pipe(audio_inputs, chunk_length_s=30, stride_length_s=[4, 2], return_timestamps="word") | |
| chunks = output["chunks"] | |
| print(chunks) | |
| # Create the video. | |
| clip = mpy.VideoClip(make_frame, duration=duration) | |
| audio_clip = mpy.AudioFileClip(audio_path).set_duration(duration) | |
| clip = clip.set_audio(audio_clip) | |
| clip.write_videofile("my_video.mp4", fps=fps, codec="libx264", audio_codec="aac") | |
| return "my_video.mp4" | |
| title = "Word-level timestamps with Whisper" | |
| description = """ | |
| This demo shows Whisper <b>word-level timestamps</b> in action using Hugging Face Transformers. It creates a video showing subtitled audio with the current word highlighted. | |
| This demo uses the <b>openai/whisper-small</b> checkpoint. Since it's only a demo, the output is limited to the first 60 seconds of audio. | |
| """ | |
| article = """ | |
| <div style='margin:20px auto;'> | |
| <p>Credits:<p> | |
| <ul> | |
| <li>Shakespeare's "Henry V" speech from <a href="https://freesound.org/people/acclivity/sounds/24096/">acclivity</a> (CC BY-NC 4.0 license) | |
| <li>Lato font by Łukasz Dziedzic (licensed under Open Font License)</li> | |
| <li>Whisper model by OpenAI</li> | |
| </ul> | |
| </div> | |
| """ | |
| examples = [ | |
| "examples/henry5.wav", | |
| ] | |
| gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Audio(label="Upload Audio", source="upload", type="filepath"), | |
| ], | |
| outputs=[ | |
| gr.Video(label="Output Video"), | |
| ], | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| ).launch() | |