Limit audio files to 120s

app.py

@@ -4,9 +4,14 @@ import gradio as gr
 from utils import write_vtt
 import whisper
 
+import ffmpeg
+
 #import os
 #os.system("pip install git+https://github.com/openai/whisper.git")
 
+# Limitations (set to -1 to disable)
+INPUT_AUDIO_MAX_DURATION = 60 # seconds
+
 LANGUAGES = [
     "English",
     "Chinese",
@@ -116,6 +121,13 @@ def greet(modelName, languageName, uploadFile, microphoneData, task):
     selectedLanguage = languageName.lower() if len(languageName) > 0 else None
     selectedModel = modelName if modelName is not None else "base"
 
+    if INPUT_AUDIO_MAX_DURATION > 0:
+        # Calculate audio length
+        audioDuration = ffmpeg.probe(source)["format"]["duration"]
+        
+        if float(audioDuration) > INPUT_AUDIO_MAX_DURATION:
+            return ("[ERROR]: Maximum audio file length is " + str(INPUT_AUDIO_MAX_DURATION) + "s, file was " + str(audioDuration) + "s"), "[ERROR]"
+
     model = model_cache.get(selectedModel, None)
     
     if not model:
@@ -130,7 +142,14 @@ def greet(modelName, languageName, uploadFile, microphoneData, task):
 
     return result["text"], segmentStream.read()
 
-demo = gr.Interface(fn=greet, description="Whisper is a general-purpose speech recognition model. It is trained on a large dataset of diverse audio and is also a multi-task model that can perform multilingual speech recognition as well as speech translation and language identification.", inputs=[
+ui_description = "Whisper is a general-purpose speech recognition model. It is trained on a large dataset of diverse " 
+ui_description += " audio and is also a multi-task model that can perform multilingual speech recognition "
+ui_description += " as well as speech translation and language identification. "
+
+if INPUT_AUDIO_MAX_DURATION > 0:
+    ui_description += "\n\n" + "Max audio file length: " + str(INPUT_AUDIO_MAX_DURATION) + " s"
+
+demo = gr.Interface(fn=greet, description=ui_description, inputs=[
     gr.Dropdown(choices=["tiny", "base", "small", "medium", "large"], value="medium", label="Model"),
     gr.Dropdown(choices=sorted(LANGUAGES), label="Language"),
     gr.Audio(source="upload", type="filepath", label="Upload Audio"), 

requirements.txt

@@ -1,2 +1,3 @@
 git+https://github.com/openai/whisper.git
-transformers
+transformers
+ffmpeg-python==0.2.0