Added captions API, that receives a URL and both transcribes AND translates it.

app.py

@@ -137,7 +137,7 @@ with gr.Blocks(css='@import "file=static/css/main.css";', theme='darkpeach', tit
   init_video.change(fn=init_video_manual_upload, inputs=[url_input, init_video], outputs=[])
 
   # Render imported buttons for API bindings
-  render_api_elements(url_input,download_status, output_text, sub_video)
+  render_api_elements(url_input,download_status, output_text, sub_video, output_file)
 
 queue_placeholder = demo.queue()
 

download.py

@@ -13,7 +13,7 @@ import argparse
 import whisper
 from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
 import ffmpeg
-from utils.subs import bake_subs
+from utils.subs import bake_subs, get_srt
 from utils.utils import get_args
 
 original_dir = os.getcwd()
@@ -106,6 +106,54 @@ def download_generator(url, translate_action=True, source_language='Autodetect',
     yield {"message": f"{e}"}
 
 
+def caption_generator(tweet_url, language="Autodetect", model_size=model_size):
+  # Download the file
+
+  try:
+    print(f"Downloading {tweet_url} ")
+    meta = check_download(tweet_url)
+    tempdir = output_dir / f"{meta['id']}"
+    print(f"Downloaded {meta['id']}.mp3 from {meta['uploader_id']} and url {meta['webpage_url']}")
+  except Exception as e:
+    print(f"Could not download file: {e}")
+    raise
+
+  try:
+    print(f"Starting audio only download with URL {tweet_url}, this may take a while")
+    meta, video, audio = download(tweet_url, tempdir, keepVideo=False)
+    print(f"Downloaded video and extracted audio")
+  except Exception as e:
+    print(f"Could not download file: {e}")
+    raise
+
+  # Run whisper on the audio with language unless auto
+  try:
+    print(f"Starting whisper transcribe with {meta['id']}.mp3")
+    transcribe_whisper_result = transcribe(audio, translate_action=False, language=language, override_model_size=model_size)
+    translate_whisper_result = transcribe(audio, translate_action=True, language=language, override_model_size=model_size)
+    srt = get_srt(transcribe_whisper_result["segments"])
+    en_srt = get_srt(translate_whisper_result["segments"])
+
+    print(f"Transcribe successful!")
+  except Exception as e:
+    print(f"Could not transcribe file: {e}")
+    return
+
+  return_dict = {
+    "detected_language": LANGUAGES[transcribe_whisper_result["language"]],
+    "requested_language": language,
+    "text": transcribe_whisper_result["text"],
+    "en_text": translate_whisper_result["text"],
+    "srt": srt,
+    "en_srt": en_srt,
+    "meta": meta,
+  }
+  return return_dict
+
+
+  # Run whisper with translation task enabled (and save to different srt file)
+  # Call anvil background task with both files, and both the plain texts
+
 def progress_hook(d):
   if d['status'] == 'downloading':
     print("downloading " + str(round(float(d['downloaded_bytes']) / float(d['total_bytes']) * 100, 1)) + "%")
@@ -115,11 +163,11 @@ def progress_hook(d):
     print(filename)
     yield f"Downloaded {filename}"
 
-def download(url, tempdir):
+def download(url, tempdir, format="bestvideo[ext=mp4]+bestaudio/best", verbose=False, keepVideo=True):
   try:
     ydl_opts = {
-      "format": "bestvideo[ext=mp4]+bestaudio/best",
-      "keepvideo": True,
+      "format": format,
+      "keepvideo": keepVideo,
       'postprocessors': [{
         'key': 'FFmpegExtractAudio',
         'preferredcodec': 'mp3',
@@ -128,7 +176,7 @@ def download(url, tempdir):
       "skip_download": False,
       "outtmpl": f"{tempdir}/%(id)s.%(ext)s",
       "noplaylist": True,
-      "verbose": False,
+      "verbose": verbose,
       "quiet": True,
       "progress_hooks": [progress_hook],
 
@@ -141,10 +189,13 @@ def download(url, tempdir):
   except DownloadError as e:
     raise e
   else:
-    video = tempdir / f"{meta['id']}.{meta['ext']}"
     audio = tempdir / f"{meta['id']}.mp3"
-    print(str(video.resolve()))
-    return meta, str(video.resolve()), str(audio.resolve())
+    if (keepVideo):
+      video = tempdir / f"{meta['id']}.{meta['ext']}"
+      return meta, str(video.resolve()), str(audio.resolve())
+    else:
+      return meta, None, str(audio.resolve())
+
 
 def check_download(url):
   ydl_opts = {
@@ -164,22 +215,24 @@ def check_download(url):
   else:
     return meta
 
-def transcribe(audio, translate_action=True, language='Autodetect'):
+def transcribe(audio, translate_action=True, language='Autodetect', override_model_size=''):
   task = "translate" if translate_action else "transcribe"
-  print(f'Starting {task} with whisper size {model_size}')
+  model_size_to_load = override_model_size if override_model_size else model_size
+  print(f'Starting {task} with whisper size {model_size_to_load} on {audio}')
   global model
-  if not preload_model:
-    model = whisper.load_model(model_size)
+  if not preload_model or model_size != override_model_size:
+    model = whisper.load_model(model_size_to_load)
+
   props = {
     "task": task,
   }
+
   if language != 'Autodetect':
     props["language"] = TO_LANGUAGE_CODE[language.lower()]
 
-  output = model.transcribe(audio, task=task)
+  output = model.transcribe(audio, verbose=True, **props)
 
-  output["language"] = LANGUAGES[output["language"]]
-  output['segments'] = [{"id": 0, "seek": 0, "start": 0.0, "end": 3, "text": " [AI transcription]"}] + output['segments']
+  output['segments'] = output['segments']
   print(f'Finished transcribe from {output["language"]}', output["text"])
   return output
 

requirements.txt

@@ -4,4 +4,5 @@ anvil-uplink==0.4.0
 gradio==3.4.0
 python-dotenv==0.21.0
 aiohttp==3.8.3
-aiohttp-requests==0.1.3
+aiohttp-requests==0.1.3
+fsspec=2022.8.2

static/css/main.css

@@ -93,5 +93,5 @@ background: transparent
 }
 
 footer{
-    display: none !important;
+    /*display: none !important;*/
 }

utils/apis.py

@@ -11,10 +11,11 @@ import anvil.media
 import dotenv
 import gradio as gr
 import requests
-from download import download_generator
-
+from download import download_generator, caption_generator
 
 dotenv.load_dotenv()
+
+
 @anvil.server.callable
 def call_gradio_api(api_name='test_api', data=()):
   port = os.environ.get('SERVER_PORT', 8111)
@@ -62,7 +63,19 @@ def test_api(url=''):
   # TODO: add an anvil server pingback to show we completed the queue operation
   return f"I've slept for 15 seconds and now I'm done. "
 
-def render_api_elements(url_input, download_status, output_text, sub_video):
+#TODO: add telegram error handler here
+def caption(tweet_url="", language="Autodetect", override_model_size=""):
+  """
+  :param media_id: The twitter media ID object
+  :param user_id_str: The twitter user ID string
+  :param tweet_url: tweet URL can potentially not exist in the future, so we can upload on behalf of the user
+  :return:
+  """
+  response = caption_generator(tweet_url, language, override_model_size)
+  return json.dumps(response)
+
+
+def render_api_elements(url_input, download_status, output_text, sub_video, output_file):
   with gr.Group(elem_id='fake_ass_group') as api_buttons:
     # This is a hack to get APIs registered with the blocks interface
     translate_result = gr.Textbox(visible=False)
@@ -75,6 +88,21 @@ def render_api_elements(url_input, download_status, output_text, sub_video):
 
     gr.Button("remote_download", visible=False)\
       .click(api_name='remote_download', queue=True, fn=remote_download, inputs=[url_input], outputs=[download_status, output_text, translate_result, translate_language])
+
+    # creating fake elements just make gradio, cause I can't define an API signature like a sane person
+
+    gr.Button("caption", visible=False)\
+      .click(api_name='caption',
+             queue=True,
+             fn=caption,
+             inputs=[
+                     gr.Text(label='tweet_url'),
+                     gr.Text(label='language (optional)'),
+                     gr.Dropdown(label='Model Size', choices=['base', 'tiny', 'small', 'medium', 'large']),
+                     ],
+             outputs=[
+                      gr.Text(label='response_json')
+                      ])
   return api_buttons
 
 
@@ -87,5 +115,4 @@ def cleanup_output_dir():
     if path.is_file():
       path.unlink()
     elif path.is_dir():
-      rmtree(path)
-
+      rmtree(path)

utils/subs.py

@@ -6,7 +6,6 @@ import os
 from typing import Iterator, TextIO
 
 
-
 def bake_subs(input_file, output_file, subs_file, fontsdir, translate_action):
   print(f"Baking {subs_file} into video... {input_file} -> {output_file}")
 
@@ -30,39 +29,39 @@ def bake_subs(input_file, output_file, subs_file, fontsdir, translate_action):
   fontstyle = f'Fontsize={sub_size},OutlineColour=&H40000000,BorderStyle=3,FontName={fontname},Bold=1'
   (
     ffmpeg.concat(
-        video.filter('subtitles', subs_file, fontsdir=fontfile, force_style=fontstyle),
-        audio, v=1, a=1
-      )
-      .overlay(watermark.filter('scale', iw / 3, -1), x='10', y='10')
-      .output(filename=output_file)
-      .run(quiet=True, overwrite_output=True)
+      video.filter('subtitles', subs_file, fontsdir=fontfile, force_style=fontstyle),
+      audio, v=1, a=1
+    )
+    .overlay(watermark.filter('scale', iw / 3, -1), x='10', y='10')
+    .output(filename=output_file)
+    .run(quiet=True, overwrite_output=True)
   )
 
 
 def str2bool(string):
-    str2val = {"True": True, "False": False}
-    if string in str2val:
-        return str2val[string]
-    else:
-        raise ValueError(
-            f"Expected one of {set(str2val.keys())}, got {string}")
+  str2val = {"True": True, "False": False}
+  if string in str2val:
+    return str2val[string]
+  else:
+    raise ValueError(
+      f"Expected one of {set(str2val.keys())}, got {string}")
 
 
 def format_timestamp(seconds: float, always_include_hours: bool = False):
-    assert seconds >= 0, "non-negative timestamp expected"
-    milliseconds = round(seconds * 1000.0)
+  assert seconds >= 0, "non-negative timestamp expected"
+  milliseconds = round(seconds * 1000.0)
 
-    hours = milliseconds // 3_600_000
-    milliseconds -= hours * 3_600_000
+  hours = milliseconds // 3_600_000
+  milliseconds -= hours * 3_600_000
 
-    minutes = milliseconds // 60_000
-    milliseconds -= minutes * 60_000
+  minutes = milliseconds // 60_000
+  milliseconds -= minutes * 60_000
 
-    seconds = milliseconds // 1_000
-    milliseconds -= seconds * 1_000
+  seconds = milliseconds // 1_000
+  milliseconds -= seconds * 1_000
 
-    hours_marker = f"{hours}:" if always_include_hours or hours > 0 else ""
-    return f"{hours_marker}{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_srt(transcript: Iterator[dict], file: TextIO):
@@ -77,8 +76,17 @@ def write_srt(transcript: Iterator[dict], file: TextIO):
         )
 
 
+def get_srt(transcript: Iterator[dict]):
+  srt = ''
+  for i, segment in enumerate(transcript, start=1):
+    srt += 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"
+  return srt
+
 def filename(path):
-    return os.path.splitext(os.path.basename(path))[0]
+  return os.path.splitext(os.path.basename(path))[0]
 
 
 
@@ -94,4 +102,4 @@ def filename(path):
 #   os.chdir(tempdirname)
 #   bake_subs(video_file_path, out_path, srt_path)
 #   anvil_media = anvil.media.from_file(out_path, 'video/mp4')
-#   print(anvil_media)
+#   print(anvil_media)