language change

Signed-off-by: Alex Volkov <alex.volkov@fundbox.com>

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "files.associations": {
+        "*.yaml": "home-assistant"
+    }
+}

app.py

@@ -5,6 +5,7 @@ from download import download_generator
 import anvil.media
 import os
 import dotenv
+from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
 
 from utils.apis import render_api_elements
 from utils.utils import get_args
@@ -19,108 +20,122 @@ gradio_share: bool = args.get("public")
 model_size: str = args.get("model")
 preload_model: str = args.get("preload")
 
+
+LANG_CHOICES = sorted([x.capitalize() for x in LANGUAGES.values()])
+LANG_CHOICES.insert(0, "Autodetect")
+
+
 url_input = gr.Textbox(label="Youtube/Twitter/etc video URL (supports many services)", value='https://twitter.com/starsonxh/status/1552945347194142720', lines=1, elem_id="url_input")
 # download_status = gr.Textbox(label="Status:", value='', lines=1, elem_id="download_status")
-download_status = gr.Checkbox(label="Status:", elem_id="download_status", interactive=False)
+download_status = gr.Checkbox(label="", elem_id="download_status", interactive=False)
 translate_action = gr.Checkbox(label="Auto translate to english", elem_id='translate_toggle', interactive=True, value=True)
 init_video = gr.Video(label="Upload video manually", visible=True, interactive=True, mirror_webcam=False)
 init_audio = gr.Audio(label="Downloaded audio", visible=False)
-output_text = gr.Textbox(label="Output text", lines=5, visible=False, max_lines=10, interactive=True)
+output_text = gr.Textbox(label="Output text", lines=5, visible=False, max_lines=10, interactive=True, elem_id="output_text")
 sub_video = gr.Video(label="Subbed video", visible=False, mirror_webcam=False)
 
 
-
-def predownload(url, translate_action):
-  for response in download_generator(url, translate_action):
+def predownload(url, translate_action, source_language):
+  files = []
+  for response in download_generator(url, translate_action, source_language):
     updates_object = {}
-    updates_object[download_status] = gr.update(label=f"STATUS: {response.get('message')}")
+    updates_object[download_status] = gr.update(label=f"{response.get('message')}")
     meta = response.get('meta')
+
     if 'video' in response:
       updates_object[init_video] = gr.update(visible=True, value=response["video"],
                                              label=f"Init Video: {meta['id']}.{meta['ext']}")
       updates_object[init_audio] = gr.update(visible=True, value=response["audio"],
                                              label=f"Extracted audio : {meta['id']}.mp3")
+      files.append(response["video"])
+      files.append(response["audio"])
     if 'whisper_result' in response:
       updates_object[output_text] = gr.update(value=response['whisper_result'].get('srt'), visible=True,
-                                              label=f"Subtitles translated from {response['whisper_result'].get('language')} (detected language)")
+                                              label=f"Subtitles transcribed from {response['whisper_result'].get('language')} (detected language)")
+    if 'srt_path' in response:
+      files.append(response["srt_path"])
+
     if 'sub_video' in response:
       updates_object[sub_video] = gr.update(visible=True, value=response["sub_video"],
                                             label=f"Subbed video: {meta['id']}_translated.mp4")
+      files.append(response["sub_video"])
 
+    updates_object[output_file] = gr.update(value=files, visible=len(files) > 0, label=f"Output Files")
     yield updates_object
 
-def rebake(output_text, sub_video):
-  raise NotImplementedError("Rebake not implemented yet")
+def correct_subtitles(url, output_text):
+  for response in download_generator(url, corrected_subtitles=output_text):
+    updates_object = {}
+    updates_object[download_status] = gr.update(label=f"STATUS: {response.get('message')}")
+    if 'sub_video' in response:
+      updates_object[sub_video] = gr.update(visible=True, value=response["sub_video"],
+                                            label=f"Corrected subtitles")
+    yield updates_object
 
 
 subtitled_video = False
-css = """
-#submit{
-  position: absolute;
-  flex:0 !important;
-  width: 120px;
-  right: 13px;
-  top: 40px;
-}
-#url_input{
-  font-size: 40px !important;
-}
-#download_status{
-  font-size: 40px !important;
-  }
-
-#input_row{
-  position: relative;
-}
-#url_input_group .gr-form:nth-child(2){
-  position:relative
-}
-#translate_toggle{
-position: absolute;
-right: 0;
-width: auto;
-flex: none;
-}
-
-#fake_ass_group{
-  display:none;
-  visibility: hidden;
-  position:absolute;
-  pointer-events: none;
-}
-    
-"""
-with gr.Blocks(css=css+"") as demo:
-  gr.Markdown('# Vid Translator 0.1 - get english subtitles for videos in any language')
-  gr.Markdown('### Link to a tweet, youtube or other video and get a translated video with @openAi #whisper, built by [@altryne](https://twitter.com/altryne/)')
-  gr.Markdown('### This is used as the backend for [@vidtranslator](https://twitter.com/vidtranslator/)')
+
+with gr.Blocks(css='@import "file=static/css/main.css";', theme='darkpeach', title='Vid Translator Studio') as demo:
+  gr.HTML('<h1 class="main-title">VidTranslator Studio 0.1</h1>')
+  gr.HTML("<h2 class='secondary'>Automatic social media video translation from 99 languages</h2>")
+
   with gr.Row(elem_id="input_row"):
     with gr.Group() as group:
       url_input.render()
       action_btn = gr.Button(elem_id='submit', variant='primary', value="Translate")
-      download_status.render()
-      translate_action.render()
+      gr.StatusTracker()
+      with gr.Row(elem_id="second_row"):
+        source_language = gr.Dropdown(choices=LANG_CHOICES,
+                                      label="Source Language",
+                                      value='Autodetect',
+                                      interactive=True, elem_id="source_language")
+        download_status.render()
+        translate_action.render()
+
   with gr.Row():
     with gr.Column():
-
       init_video.render()
       init_audio.render()
+      with gr.Row():
+          output_file = gr.Files(label="Output Files", visible=False)
+
     with gr.Column():
-      with gr.Group():
-        sub_video.render()
-        output_text.render()
-        gr.Button("Download srt file")
-        rebake = gr.Button("Edit subtitles on video")
-
-  outputs = [download_status, init_video, init_audio, output_text, sub_video]
-  inputs = [url_input, translate_action]
+      output_text.render()
+      correct_btn = gr.Button("Correct subtitles")
+
+    with gr.Column():
+      sub_video.render()
+
+
+  outputs = [download_status, init_video, init_audio, output_text, sub_video, output_file ]
+  inputs = [url_input, translate_action, source_language]
   action_btn.click(fn=predownload, inputs=inputs, outputs=outputs, api_name='predownload')
   url_input.submit(fn=predownload, inputs=inputs, outputs=outputs)
-  rebake.click(fn=rebake, inputs=[output_text, sub_video], outputs=[download_status, output_text, sub_video])
+
+  correct_btn.click(fn=correct_subtitles, inputs=[url_input, output_text], outputs=[download_status, output_text, sub_video])
 
   translate_action.change(fn=lambda x: {action_btn: gr.update(value=f"Translate" if x else "Transcribe")},
                           inputs=[translate_action], outputs=[action_btn])
 
+  gr.HTML("""<div class='footer'>
+    <div class="relative">
+  <div class="absolute inset-0 flex items-center" aria-hidden="true">
+    <div class="w-full border-t border-gray-300"></div>
+  </div>
+  <div class="relative flex justify-center">
+    <span class="bg-white px-2 text-sm text-gray-500">Continue</span>
+  </div>
+</div>
+
+     <a href='https://twitter.com/altryne/'><img src='https://img.shields.io/badge/%40altryne-follow-green' alt=''></a>
+     </div>""")
+
+  def init_video_manual_upload(url, init_video):
+    print(url)
+    print(init_video)
+
+  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)
 

download.py

@@ -11,7 +11,7 @@ import tempfile
 import json
 import argparse
 import whisper
-from whisper.tokenizer import LANGUAGES
+from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
 import ffmpeg
 from utils.subs import bake_subs
 from utils.utils import get_args
@@ -27,9 +27,9 @@ if preload_model:
   print("Preloading model")
   model = whisper.load_model(model_size)
 
-def download_generator(url, translate_action=True):
+def download_generator(url, translate_action=True, source_language='Autodetect', corrected_subtitles=None):
 
-  ### Step 1 : check if video is available
+  # Step 1 : check if video is available
   yield {"message": f"Checking {url} for videos"}
   try:
     meta = check_download(url)
@@ -43,32 +43,46 @@ def download_generator(url, translate_action=True):
     yield {"message": f"{e}"}
     return
 
-  ### Step 2 : Download video and extract audio
-  yield {"message": f"Starting download with URL {url}, this may take a while"}
-
+  # Step 2 : Download video and extract audio
   try:
-    meta, video, audio = download(url, tempdir)
+    # check if we already have the folder and the main files
+    if(tempdir.is_dir() and (tempdir/f"{meta['id']}.{meta['ext']}").is_file() and (tempdir/f"{meta['id']}.mp3").is_file()):
+      yield {"message": f"Using cached files"}
+      video = str((tempdir/f"{meta['id']}.{meta['ext']}").resolve())
+      audio = str((tempdir/f"{meta['id']}.mp3").resolve())
+    else:
+      yield {"message": f"Starting download with URL {url}, this may take a while"}
+      meta, video, audio = download(url, tempdir)
     yield {"message": f"Downloaded video and extracted audio", "video": video, "audio": audio, "meta": meta}
+
   except Exception as e:
     os.chdir(original_dir)
     yield {"message": f"{e}"}
     raise e
 
-  ### Step 3 : Transcribe with whisper
-  yield {"message": f"[PLEASE WAIT] Starting whisper transcribe with {meta['id']}.mp3"}
-  try:
-    whisper_result = transcribe(audio, translate_action)
-    srt_path = tempdir / f"{meta['id']}.srt"
-    with open(srt_path, "w", encoding="utf-8") as srt:
-      write_srt(whisper_result["segments"], file=srt)
+  srt_path = tempdir / f"{meta['id']}.srt"
 
-    whisper_result["srt"] = Path(srt_path).read_text()
-    yield {"message": f"Transcribe successful", "whisper_result": whisper_result, "meta": meta}
-  except Exception as e:
-    os.chdir(original_dir)
-    yield {"message": f"{e}"}
-    raise e
-    #TODO: add return here so users can continue after editing subtitles
+  if not corrected_subtitles:
+    ### Step 3 : Transcribe with whisper
+    yield {"message": f"[PLEASE WAIT] Starting whisper transcribe with {meta['id']}.mp3"}
+    try:
+      whisper_result = transcribe(audio, translate_action, source_language)
+
+      with open(srt_path, "w", encoding="utf-8") as srt:
+        write_srt(whisper_result["segments"], file=srt)
+
+      whisper_result["srt"] = Path(srt_path).read_text()
+      yield {"message": f"Transcribe successful", "whisper_result": whisper_result, "meta": meta, "srt_path": srt_path}
+    except Exception as e:
+      os.chdir(original_dir)
+      yield {"message": f"{e}"}
+      raise e
+  else:
+    ### step 3.5 : use corrected subtitles
+    yield {"message": f"Using corrected subtitles"}
+    with open(srt_path, "w", encoding="utf-8") as srt:
+      srt.write(corrected_subtitles)
+    yield {"message": f"Transcribe successful", "srt_path": srt_path, "meta": meta}
 
   ### Step 4 : Bake subtitles into video with ffmpeg
   yield {"message": f"[PLEASE WAIT] baking subtitles into video"}
@@ -150,14 +164,22 @@ def check_download(url):
   else:
     return meta
 
-def transcribe(audio, translate_action=True):
-  print(f'Starting ' + "translate" if translate_action else "transcribe")
+def transcribe(audio, translate_action=True, language='Autodetect'):
+  task = "translate" if translate_action else "transcribe"
+  print(f'Starting {task} with whisper size {model_size}')
   global model
   if not preload_model:
     model = whisper.load_model(model_size)
-  output = model.transcribe(audio, task="translate" if translate_action else "transcribe")
+  props = {
+    "task": task,
+  }
+  if language != 'Autodetect':
+    props["language"] = TO_LANGUAGE_CODE[language.lower()]
+
+  output = model.transcribe(audio, task=task)
+
   output["language"] = LANGUAGES[output["language"]]
-  output['segments'] = [{"id": 0, "seek": 0, "start": 0.0, "end": 3, "text": " [AI translation by @vidtranslator]"}] + output['segments']
+  output['segments'] = [{"id": 0, "seek": 0, "start": 0.0, "end": 3, "text": " [AI transcription]"}] + output['segments']
   print(f'Finished transcribe from {output["language"]}', output["text"])
   return output
 

environment.yml

@@ -0,0 +1,8 @@
+name: vidnightly
+channels:
+  - defaults
+dependencies:
+  - pytorch
+  - torchvision
+  - torchaudio
+prefix: /opt/homebrew/Caskroom/miniconda/base/envs/vidnightly

static/css/main.css

@@ -0,0 +1,97 @@
+@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@700&display=swap');
+
+#download_status input[type=checkbox] {
+  display: none;
+}
+#download_status label{
+    height: 100%;
+}
+.main-title{
+    color: #FF7A7A;
+    font-family: 'Poppins', sans-serif;
+    font-size: 4.5em;
+    line-height: 1.25;
+    font-weight: 700;
+    background-image: linear-gradient(45deg, #695EE6 0%, #FF7A7A 85%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+}
+.secondary{
+    color: #FF7A7A;
+    font-family: 'Poppins', sans-serif;
+    font-size: 2em;
+    line-height: 1.25;
+    font-weight: 700;
+    background-image: linear-gradient(45deg, #695EE6 0%, #FF7A7A 85%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+}
+#submit{
+  position: absolute;
+  flex:0 !important;
+  width: 120px;
+  right: 13px;
+  top: 40px;
+  border-radius: 0 5px 5px 5px !important;
+}
+#url_input{
+  font-size: 20px !important;
+}
+#download_status label{
+  font-size: 18px !important;
+}
+#second_row>.gr-form{
+  border-top-left-radius: 0px !important;
+  border-top-right-radius: 0px !important;
+}
+
+#input_row{
+  position: relative;
+}
+#url_input_group .gr-form:nth-child(2){
+  position:relative
+}
+#url_input textarea{
+  font-size: 20px !important;
+  }
+
+#source_language{
+flex-grow: revert
+}
+
+#translate_toggle{
+flex-grow: revert;
+min-width: 200px;
+}
+#translate_toggle label{
+    height: 100%;
+}
+
+/*
+#translate_toggle{
+position: absolute;
+right: 0;
+width: auto;
+flex: none;
+background: transparent
+}
+*/
+.wrap.absolute{
+    position: relative !important;
+    opacity: 100% !important;
+}
+
+#fake_ass_group{
+  display:none;
+  visibility: hidden;
+  position:absolute;
+  pointer-events: none;
+}
+#output_text label{
+  font-size: 20px !important;
+  white-space: pre-wrap;
+}
+
+footer{
+    display: none !important;
+}