Add support for downloading the results.

app.py

@@ -1,7 +1,11 @@
 from io import StringIO
+import os
+import tempfile
+
+from typing import Iterator
 import gradio as gr
 
-from utils import write_vtt
+from utils import slugify, write_srt, write_vtt
 import whisper
 
 import ffmpeg
@@ -40,6 +44,8 @@ class UI:
 
     def transcribeFile(self, modelName, languageName, uploadFile, microphoneData, task):
         source = uploadFile if uploadFile is not None else microphoneData
+        sourceName = os.path.basename(source)
+
         selectedLanguage = languageName.lower() if len(languageName) > 0 else None
         selectedModel = modelName if modelName is not None else "base"
 
@@ -56,14 +62,43 @@ class UI:
             model = whisper.load_model(selectedModel)
             model_cache[selectedModel] = model
 
+        # The results
         result = model.transcribe(source, language=selectedLanguage, task=task)
 
-        segmentStream = StringIO()
-        write_vtt(result["segments"], file=segmentStream)
-        segmentStream.seek(0)
+        text = result["text"]
+        vtt = getSubs(result["segments"], "vtt")
+        srt = getSubs(result["segments"], "srt")
+
+        # Files that can be downloaded
+        downloadDirectory = tempfile.mkdtemp()
+        filePrefix = slugify(sourceName, allow_unicode=True)
+
+        download = []
+        download.append(createFile(srt, downloadDirectory, filePrefix + "-subs.srt"));
+        download.append(createFile(vtt, downloadDirectory, filePrefix + "-subs.vtt"));
+        download.append(createFile(text, downloadDirectory, filePrefix + "-transcript.txt"));
+
+        return text, vtt, download
+
+def createFile(text: str, directory: str, fileName: str) -> str:
+    # Write the text to a file
+    with open(os.path.join(directory, fileName), 'w+', encoding="utf-8") as file:
+        file.write(text)
+
+    return file.name
+
+def getSubs(segments: Iterator[dict], format: str) -> str:
+    segmentStream = StringIO()
 
-        return result["text"], segmentStream.read()
+    if format == 'vtt':
+        write_vtt(segments, file=segmentStream)
+    elif format == 'srt':
+        write_srt(segments, file=segmentStream)
+    else:
+        raise Exception("Unknown format " + format)
 
+    segmentStream.seek(0)
+    return segmentStream.read()
 
 def createUi(inputAudioMaxDuration, share=False):
     ui = UI(inputAudioMaxDuration)
@@ -81,7 +116,11 @@ def createUi(inputAudioMaxDuration, share=False):
         gr.Audio(source="upload", type="filepath", label="Upload Audio"), 
         gr.Audio(source="microphone", type="filepath", label="Microphone Input"),
         gr.Dropdown(choices=["transcribe", "translate"], label="Task"),
-    ], outputs=[gr.Text(label="Transcription"), gr.Text(label="Segments")])
+    ], outputs=[
+        gr.Text(label="Transcription"), 
+        gr.Text(label="Segments"),
+        gr.File(label="Download")
+    ])
 
     demo.launch(share=share)   
 

utils.py

@@ -1,3 +1,6 @@
+import unicodedata
+import re
+
 import zlib
 from typing import Iterator, TextIO
 
@@ -27,7 +30,7 @@ def compression_ratio(text) -> float:
     return len(text) / len(zlib.compress(text.encode("utf-8")))
 
 
-def format_timestamp(seconds: float):
+def format_timestamp(seconds: float, always_include_hours: bool = False):
     assert seconds >= 0, "non-negative timestamp expected"
     milliseconds = round(seconds * 1000.0)
 
@@ -40,7 +43,13 @@ def format_timestamp(seconds: float):
     seconds = milliseconds // 1_000
     milliseconds -= seconds * 1_000
 
-    return (f"{hours}:" if hours > 0 else "") + f"{minutes:02d}:{seconds:02d}.{milliseconds:03d}"
+    hours_marker = f"{hours}:" if always_include_hours or hours > 0 else ""
+    return f"{hours_marker}{minutes:02d}:{seconds:02d}.{milliseconds:03d}"
+
+
+def write_txt(transcript: Iterator[dict], file: TextIO):
+    for segment in transcript:
+        print(segment['text'].strip(), file=file, flush=True)
 
 
 def write_vtt(transcript: Iterator[dict], file: TextIO):
@@ -52,3 +61,43 @@ def write_vtt(transcript: Iterator[dict], file: TextIO):
             file=file,
             flush=True,
         )
+
+
+def write_srt(transcript: Iterator[dict], file: TextIO):
+    """
+    Write a transcript to a file in SRT format.
+    Example usage:
+        from pathlib import Path
+        from whisper.utils import write_srt
+        result = transcribe(model, audio_path, temperature=temperature, **args)
+        # save SRT
+        audio_basename = Path(audio_path).stem
+        with open(Path(output_dir) / (audio_basename + ".srt"), "w", encoding="utf-8") as srt:
+            write_srt(result["segments"], file=srt)
+    """
+    for i, segment in enumerate(transcript, start=1):
+        # write srt lines
+        print(
+            f"{i}\n"
+            f"{format_timestamp(segment['start'], always_include_hours=True)} --> "
+            f"{format_timestamp(segment['end'], always_include_hours=True)}\n"
+            f"{segment['text'].strip().replace('-->', '->')}\n",
+            file=file,
+            flush=True,
+        )
+
+def slugify(value, allow_unicode=False):
+    """
+    Taken from https://github.com/django/django/blob/master/django/utils/text.py
+    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated
+    dashes to single dashes. Remove characters that aren't alphanumerics,
+    underscores, or hyphens. Convert to lowercase. Also strip leading and
+    trailing whitespace, dashes, and underscores.
+    """
+    value = str(value)
+    if allow_unicode:
+        value = unicodedata.normalize('NFKC', value)
+    else:
+        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
+    value = re.sub(r'[^\w\s-]', '', value.lower())
+    return re.sub(r'[-\s]+', '-', value).strip('-_')