app.py

import gradio
import gradio as gr

from download import download_generator, user_uploaded_video_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

dotenv.load_dotenv()

anvil.server.connect(os.environ.get('ANVIL_UPLINK_KEY'))
queue_placeholder = None

args = get_args()
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")

VIDEO_HTML = """
<video
  class="video-js"
  controls
  preload="auto"
  width="640"
  height="264"
  >
  <source src="{src}" type="video/mp4">
  <track kind="captions" src="{en_vtt}" srclang="en" label="English" default>
</video>
"""

url_input = gr.Textbox(label="Youtube/Twitter/etc video URL (supports many services)", lines=1, elem_id="url_input")
# download_status = gr.Textbox(label="Status:", value='', lines=1, elem_id="download_status")
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, elem_id="output_text")
output_text_2 = gr.Textbox(label="Output text 2", 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)
sub_video_html = gr.HTML(value=f"<div> Please wait for video to load </div>")

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"{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 transcribed from {response['whisper_result'].get('language')} (detected language)")
    if 'srt_path' in response:
      files.append(response["srt_path"])
    if 'vtt_path' in response:
      files.append(response["vtt_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")
      updates_object[sub_video_html] = gr.update(value=VIDEO_HTML.format(src=f"file={response['sub_video']}", en_vtt=f"file={response['vtt_path']}") )
      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 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

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")
      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():
      output_text.render()
      correct_btn = gr.Button("Correct subtitles")

    with gr.Column():
      sub_video.render()
      sub_video_html.render()


  outputs = [download_status, init_video, init_audio, output_text, sub_video, output_file, sub_video_html]
  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)

  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])
  examples = gr.Examples([["https://twitter.com/starsonxh/status/1552945347194142720", "Adam"], ["https://twitter.com/starsonxh/status/1552945347194142720", "Eve"]], [url_input, output_text] )
  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):
    if url:
      return False
    files = []
    for response in user_uploaded_video_generator(init_video):
      updates_object = {}
      updates_object[download_status] = gr.update(label=f"{response.get('message')}")



      if 'audio' in response:
        updates_object[init_audio] = gr.update(visible=True, value=response["audio"],
                                               label=f"Extracted audio")
        files.append(response["audio"])
        files.append(response["video"])


      if 'srt_path' in response:
        updates_object[output_text] = gr.update(value=response['srt_path'], visible=True)
        files.append(response["srt_path"])
        updates_object[sub_video_html] = gr.update(value=VIDEO_HTML % f"file={response['sub_video']}")

      if 'vtt_path' in response:
        updates_object[output_text_2] = gr.update(value=response['vtt_path'], visible=True)
        files.append(response["vtt_path"])
        updates_object[sub_video_html] = gr.update(value=VIDEO_HTML.format(src=f"file={response['sub_video']}", en_vtt=f"file={response['vtt_path']}"))
      #
      # updates_object[output_file] = gr.update(value=files, visible=len(files) > 0, label=f"Output Files")

      yield updates_object


  init_video.change(fn=init_video_manual_upload,
                    inputs=[url_input, init_video],
                    outputs=[download_status, init_audio, sub_video_html, output_file])

  # Render imported buttons for API bindings
  render_api_elements(url_input,download_status, output_text, sub_video, output_file)

queue_placeholder = demo.queue()



if __name__ == "__main__":
  gradio.close_all()
  port = os.environ.get('SERVER_PORT', 8111)
  demo.launch(show_error=True, debug=True, share=gradio_share,server_port=int(port), favicon_path='fonts/icon.png')