Good trascription and translation, but text2speech doesn't works good

lang_list.py

@@ -1,360 +1,175 @@
-# Language dict
-language_code_to_name = {
-    "afr": "Afrikaans",
-    "amh": "Amharic",
-    "arb": "Modern Standard Arabic",
-    "ary": "Moroccan Arabic",
-    "arz": "Egyptian Arabic",
-    "asm": "Assamese",
-    "ast": "Asturian",
-    "azj": "North Azerbaijani",
-    "bel": "Belarusian",
-    "ben": "Bengali",
-    "bos": "Bosnian",
-    "bul": "Bulgarian",
-    "cat": "Catalan",
-    "ceb": "Cebuano",
-    "ces": "Czech",
-    "ckb": "Central Kurdish",
-    "cmn": "Mandarin Chinese",
-    "cym": "Welsh",
-    "dan": "Danish",
-    "deu": "German",
-    "ell": "Greek",
-    "eng": "English",
-    "est": "Estonian",
-    "eus": "Basque",
-    "fin": "Finnish",
-    "fra": "French",
-    "gaz": "West Central Oromo",
-    "gle": "Irish",
-    "glg": "Galician",
-    "guj": "Gujarati",
-    "heb": "Hebrew",
-    "hin": "Hindi",
-    "hrv": "Croatian",
-    "hun": "Hungarian",
-    "hye": "Armenian",
-    "ibo": "Igbo",
-    "ind": "Indonesian",
-    "isl": "Icelandic",
-    "ita": "Italian",
-    "jav": "Javanese",
-    "jpn": "Japanese",
-    "kam": "Kamba",
-    "kan": "Kannada",
-    "kat": "Georgian",
-    "kaz": "Kazakh",
-    "kea": "Kabuverdianu",
-    "khk": "Halh Mongolian",
-    "khm": "Khmer",
-    "kir": "Kyrgyz",
-    "kor": "Korean",
-    "lao": "Lao",
-    "lit": "Lithuanian",
-    "ltz": "Luxembourgish",
-    "lug": "Ganda",
-    "luo": "Luo",
-    "lvs": "Standard Latvian",
-    "mai": "Maithili",
-    "mal": "Malayalam",
-    "mar": "Marathi",
-    "mkd": "Macedonian",
-    "mlt": "Maltese",
-    "mni": "Meitei",
-    "mya": "Burmese",
-    "nld": "Dutch",
-    "nno": "Norwegian Nynorsk",
-    "nob": "Norwegian Bokm\u00e5l",
-    "npi": "Nepali",
-    "nya": "Nyanja",
-    "oci": "Occitan",
-    "ory": "Odia",
-    "pan": "Punjabi",
-    "pbt": "Southern Pashto",
-    "pes": "Western Persian",
-    "pol": "Polish",
-    "por": "Portuguese",
-    "ron": "Romanian",
-    "rus": "Russian",
-    "slk": "Slovak",
-    "slv": "Slovenian",
-    "sna": "Shona",
-    "snd": "Sindhi",
-    "som": "Somali",
-    "spa": "Spanish",
-    "srp": "Serbian",
-    "swe": "Swedish",
-    "swh": "Swahili",
-    "tam": "Tamil",
-    "tel": "Telugu",
-    "tgk": "Tajik",
-    "tgl": "Tagalog",
-    "tha": "Thai",
-    "tur": "Turkish",
-    "ukr": "Ukrainian",
-    "urd": "Urdu",
-    "uzn": "Northern Uzbek",
-    "vie": "Vietnamese",
-    "xho": "Xhosa",
-    "yor": "Yoruba",
-    "yue": "Cantonese",
-    "zlm": "Colloquial Malay",
-    "zsm": "Standard Malay",
-    "zul": "Zulu",
+# Languages dict
+LANGUAGE_NAME_TO_CODE = {
+    "العربية": "ar_AR",
+    "Čeština": "cs_CZ",
+    "Deutsch": "de_DE",
+    "English": "en_XX",
+    "Español": "es_XX",
+    "Eesti": "et_EE",
+    "Suomi": "fi_FI",
+    "Français": "fr_XX",
+    "ગુજરાતી": "gu_IN",
+    "हिन्दी": "hi_IN",
+    "Italiano": "it_IT",
+    "日本語": "ja_XX",
+    "Қазақ": "kk_KZ",
+    "한국어": "ko_KR",
+    "Lietuvių": "lt_LT",
+    "Latviešu": "lv_LV",
+    "ဗမာ": "my_MM",
+    "नेपाली": "ne_NP",
+    "Nederlands": "nl_XX",
+    "Română": "ro_RO",
+    "Русский": "ru_RU",
+    "සිංහල": "si_LK",
+    "Türkçe": "tr_TR",
+    "Tiếng Việt": "vi_VN",
+    "中文": "zh_CN",
+    "Afrikaans": "af_ZA",
+    "Azərbaycan": "az_AZ",
+    "বাংলা": "bn_IN",
+    "فارسی": "fa_IR",
+    "עברית": "he_IL",
+    "Hrvatski": "hr_HR",
+    "Indonesia": "id_ID",
+    "ქართული": "ka_GE",
+    "ខ្មែរ": "km_KH",
+    "Македонски": "mk_MK",
+    "മലയാളം": "ml_IN",
+    "Монгол": "mn_MN",
+    "मराठी": "mr_IN",
+    "Polski": "pl_PL",
+    "پښتو": "ps_AF",
+    "Português": "pt_XX",
+    "Svenska": "sv_SE",
+    "Kiswahili": "sw_KE",
+    "தமிழ்": "ta_IN",
+    "తెలుగు": "te_IN",
+    "ไทย": "th_TH",
+    "Tagalog": "tl_XX",
+    "Українська": "uk_UA",
+    "اردو": "ur_PK",
+    "isiXhosa": "xh_ZA",
+    "Galego": "gl_ES",
+    "Slovenščina": "sl_SI"
 }
-original_language_code_to_name = {
-    "afr": "Afrikaans",
-    "amh": "አማርኛ",
-    "arb": "العربية",
-    "ary": "الدارجة المغربية",
-    "arz": "العامية المصرية",
-    "asm": "অসমীয়া",
-    "ast": "Asturianu",
-    "azj": "Azərbaycanca",
-    "bel": "Беларуская",
-    "ben": "বাংলা",
-    "bos": "Bosanski",
-    "bul": "Български",
-    "cat": "Català",
-    "ceb": "Cebuano",
-    "ces": "Čeština",
-    "ckb": "کوردی ناوەندی",
-    "cmn": "普通话",
-    "cym": "Cymraeg",
-    "dan": "Dansk",
-    "deu": "Deutsch",
-    "ell": "Ελληνικά",
-    "eng": "English",
-    "est": "Eesti",
-    "eus": "Euskara",
-    "fin": "Suomi",
-    "fra": "Français",
-    "gaz": "Afaan Oromoo",
-    "gle": "Gaeilge",
-    "glg": "Galego",
-    "guj": "ગુજરાતી",
-    "heb": "עברית",
-    "hin": "हिंदी",
-    "hrv": "Hrvatski",
-    "hun": "Magyar",
-    "hye": "Հայերեն",
-    "ibo": "Igbo",
-    "ind": "Bahasa Indonesia",
-    "isl": "Íslenska",
-    "ita": "Italiano",
-    "jav": "Basa Jawa",
-    "jpn": "日本語",
-    "kam": "Kikamba",
-    "kan": "ಕನ್ನಡ",
-    "kat": "ქართული",
-    "kaz": "Қазақ тілі",
-    "kea": "Kriolu di Kabuverdianu",
-    "khk": "Халх",
-    "khm": "ខ្មែរ",
-    "kir": "Кыргызча",
-    "kor": "한국어",
-    "lao": "ລາວ",
-    "lit": "Lietuvių",
-    "ltz": "Lëtzebuergesch",
-    "lug": "Luganda",
-    "luo": "Dholuo",
-    "lvs": "Latviešu",
-    "mai": "मैथिली",
-    "mal": "മലയാളം",
-    "mar": "मराठी",
-    "mkd": "Македонски",
-    "mlt": "Malti",
-    "mni": "মৈতৈলোন",
-    "mya": "မြန်မာ",
-    "nld": "Nederlands",
-    "nno": "Nynorsk",
-    "nob": "Bokmål",
-    "npi": "नेपाली",
-    "nya": "Chichewa",
-    "oci": "Occitan",
-    "ory": "ଓଡ଼ିଆ",
-    "pan": "ਪੰਜਾਬੀ",
-    "pbt": "پښتو",
-    "pes": "فارسی",
-    "pol": "Polski",
-    "por": "Português",
-    "ron": "Română",
-    "rus": "Русский",
-    "slk": "Slovenčina",
-    "slv": "Slovenščina",
-    "sna": "ChiShona",
-    "snd": "سنڌي",
-    "som": "Soomaali",
-    "spa": "Español",
-    "srp": "Српски",
-    "swe": "Svenska",
-    "swh": "Kiswahili",
-    "tam": "தமிழ்",
-    "tel": "తెలుగు",
-    "tgk": "Тоҷикӣ",
-    "tgl": "Tagalog",
-    "tha": "ไทย",
-    "tur": "Türkçe",
-    "ukr": "Українська",
-    "urd": "اردو",
-    "uzn": "O‘zbekcha",
-    "vie": "Tiếng Việt",
-    "xho": "IsiXhosa",
-    "yor": "Yorùbá",
-    "yue": "粤语",
-    "zlm": "Bahasa Melayu",
-    "zsm": "Bahasa Melayu",
-    "zul": "IsiZulu",
-}
-LANGUAGE_NAME_TO_CODE = {v: k for k, v in language_code_to_name.items()}
-ORIGINAL_LANGUAGE_NAME_TO_CODE = {v: k for k, v in original_language_code_to_name.items()}
-
-# Source langs: S2ST / S2TT / ASR don't need source lang
-# T2TT / T2ST use this
-text_source_language_codes = [
-    "afr",
-    "amh",
-    "arb",
-    "ary",
-    "arz",
-    "asm",
-    "azj",
-    "bel",
-    "ben",
-    "bos",
-    "bul",
-    "cat",
-    "ceb",
-    "ces",
-    "ckb",
-    "cmn",
-    "cym",
-    "dan",
-    "deu",
-    "ell",
-    "eng",
-    "est",
-    "eus",
-    "fin",
-    "fra",
-    "gaz",
-    "gle",
-    "glg",
-    "guj",
-    "heb",
-    "hin",
-    "hrv",
-    "hun",
-    "hye",
-    "ibo",
-    "ind",
-    "isl",
-    "ita",
-    "jav",
-    "jpn",
-    "kan",
-    "kat",
-    "kaz",
-    "khk",
-    "khm",
-    "kir",
-    "kor",
-    "lao",
-    "lit",
-    "lug",
-    "luo",
-    "lvs",
-    "mai",
-    "mal",
-    "mar",
-    "mkd",
-    "mlt",
-    "mni",
-    "mya",
-    "nld",
-    "nno",
-    "nob",
-    "npi",
-    "nya",
-    "ory",
-    "pan",
-    "pbt",
-    "pes",
-    "pol",
-    "por",
-    "ron",
-    "rus",
-    "slk",
-    "slv",
-    "sna",
-    "snd",
-    "som",
-    "spa",
-    "srp",
-    "swe",
-    "swh",
-    "tam",
-    "tel",
-    "tgk",
-    "tgl",
-    "tha",
-    "tur",
-    "ukr",
-    "urd",
-    "uzn",
-    "vie",
-    "yor",
-    "yue",
-    "zsm",
-    "zul",
-]
-TEXT_SOURCE_LANGUAGE_NAMES = sorted([language_code_to_name[code] for code in text_source_language_codes])
 
-# Target langs:
-# S2ST / T2ST
-s2st_target_language_codes = [
-    "eng",
-    "arb",
-    "ben",
-    "cat",
-    "ces",
-    "cmn",
-    "cym",
-    "dan",
-    "deu",
-    "est",
-    "fin",
-    "fra",
-    "hin",
-    "ind",
-    "ita",
-    "jpn",
-    "kor",
-    "mlt",
-    "nld",
-    "pes",
-    "pol",
-    "por",
-    "ron",
-    "rus",
-    "slk",
-    "spa",
-    "swe",
-    "swh",
-    "tel",
-    "tgl",
-    "tha",
-    "tur",
-    "ukr",
-    "urd",
-    "uzn",
-    "vie",
-]
-S2ST_TARGET_LANGUAGE_NAMES = sorted([language_code_to_name[code] for code in s2st_target_language_codes])
-S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES = sorted([original_language_code_to_name[code] for code in s2st_target_language_codes])
+# Whisper languages dict
+WHISPER_LANGUAGES = {
+    "en": "english",
+    "zh": "chinese",
+    "de": "german",
+    "es": "spanish",
+    "ru": "russian",
+    "ko": "korean",
+    "fr": "french",
+    "ja": "japanese",
+    "pt": "portuguese",
+    "tr": "turkish",
+    "pl": "polish",
+    "ca": "catalan",
+    "nl": "dutch",
+    "ar": "arabic",
+    "sv": "swedish",
+    "it": "italian",
+    "id": "indonesian",
+    "hi": "hindi",
+    "fi": "finnish",
+    "vi": "vietnamese",
+    "he": "hebrew",
+    "uk": "ukrainian",
+    "el": "greek",
+    "ms": "malay",
+    "cs": "czech",
+    "ro": "romanian",
+    "da": "danish",
+    "hu": "hungarian",
+    "ta": "tamil",
+    "no": "norwegian",
+    "th": "thai",
+    "ur": "urdu",
+    "hr": "croatian",
+    "bg": "bulgarian",
+    "lt": "lithuanian",
+    "la": "latin",
+    "mi": "maori",
+    "ml": "malayalam",
+    "cy": "welsh",
+    "sk": "slovak",
+    "te": "telugu",
+    "fa": "persian",
+    "lv": "latvian",
+    "bn": "bengali",
+    "sr": "serbian",
+    "az": "azerbaijani",
+    "sl": "slovenian",
+    "kn": "kannada",
+    "et": "estonian",
+    "mk": "macedonian",
+    "br": "breton",
+    "eu": "basque",
+    "is": "icelandic",
+    "hy": "armenian",
+    "ne": "nepali",
+    "mn": "mongolian",
+    "bs": "bosnian",
+    "kk": "kazakh",
+    "sq": "albanian",
+    "sw": "swahili",
+    "gl": "galician",
+    "mr": "marathi",
+    "pa": "punjabi",
+    "si": "sinhala",
+    "km": "khmer",
+    "sn": "shona",
+    "yo": "yoruba",
+    "so": "somali",
+    "af": "afrikaans",
+    "oc": "occitan",
+    "ka": "georgian",
+    "be": "belarusian",
+    "tg": "tajik",
+    "sd": "sindhi",
+    "gu": "gujarati",
+    "am": "amharic",
+    "yi": "yiddish",
+    "lo": "lao",
+    "uz": "uzbek",
+    "fo": "faroese",
+    "ht": "haitian creole",
+    "ps": "pashto",
+    "tk": "turkmen",
+    "nn": "nynorsk",
+    "mt": "maltese",
+    "sa": "sanskrit",
+    "lb": "luxembourgish",
+    "my": "myanmar",
+    "bo": "tibetan",
+    "tl": "tagalog",
+    "mg": "malagasy",
+    "as": "assamese",
+    "tt": "tatar",
+    "haw": "hawaiian",
+    "ln": "lingala",
+    "ha": "hausa",
+    "ba": "bashkir",
+    "jw": "javanese",
+    "su": "sundanese",
+}
 
-# S2TT / ASR
-S2TT_TARGET_LANGUAGE_NAMES = TEXT_SOURCE_LANGUAGE_NAMES
-# T2TT
-T2TT_TARGET_LANGUAGE_NAMES = TEXT_SOURCE_LANGUAGE_NAMES
+def union_language_dict():
+    # Create a dictionary to store the language codes
+    language_dict = {}
+    # Iterate over the LANGUAGE_NAME_TO_CODE dictionary
+    for language_name, language_code in LANGUAGE_NAME_TO_CODE.items():
+        # Extract the language code (the first two characters before the underscore)
+        lang_code = language_code.split('_')[0].lower()
+        
+        # Check if the language code is present in WHISPER_LANGUAGES
+        if lang_code in WHISPER_LANGUAGES:
+            # Construct the entry for the resulting dictionary
+            language_dict[language_name] = {
+                "transcriber": lang_code,
+                "translator": language_code
+            }
+    return language_dict

requirements.txt

@@ -10,4 +10,7 @@ twitch-dl
 pytube
 pyperclip
 transformers
-git+https://github.com/openai/whisper.git
+git+https://github.com/openai/whisper.git
+sentencepiece
+protobuf
+git+https://github.com/suno-ai/bark.git

translatube.py

@@ -4,16 +4,21 @@ import urllib.parse as urlparse
 from pytube import YouTube
 import re
 import subprocess
-import time
-from lang_list import ORIGINAL_LANGUAGE_NAME_TO_CODE, S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES
-
+import torch
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast, pipeline
+from lang_list import union_language_dict
 import torch
 import whisper
+from bark import SAMPLE_RATE, generate_audio, preload_models
+from scipy.io.wavfile import write as write_wav
+import gc
 
 # get device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # device = torch.device("cpu")
-model = whisper.load_model("large-v2", device=device)
+
+# Create a dictionary to store the language codes
+language_dict = union_language_dict()
 
 YOUTUBE = "youtube"
 TWITCH = "twitch"
@@ -25,22 +30,19 @@ def copy_url_from_clipboard():
 def clear_video_url():
     visible = False
     image = gr.Image(visible=visible, scale=1)
-    source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=ORIGINAL_LANGUAGE_NAME_TO_CODE, scale=1, interactive=True)
-    target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES, scale=1, interactive=True)
-    get_audio_button = gr.Button(size="lg", value="get audio", min_width="10px", scale=0, visible=visible)
-    transcribe_audio_button = gr.Button(size="lg", value="transcribe audio", min_width="10px", scale=0, visible=visible)
+    source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=language_dict, scale=1, interactive=True)
+    target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=language_dict, scale=1, interactive=True)
+    translate_button = gr.Button(size="lg", value="translate", min_width="10px", scale=0, visible=visible)
     original_audio = gr.Audio(label="Original audio", elem_id="original_audio", visible=visible, interactive=False)
     original_audio_transcribed = gr.Textbox(label="Original audio transcribed", elem_id="original_audio_transcribed", interactive=False, visible=visible)
     original_audio_translated = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", interactive=False, visible=visible)
     translated_audio = gr.Audio(label="Translated audio", elem_id="translated_audio", visible=visible)
-    transcribe_audio_button = gr.Button(size="lg", value="transcribe audio", min_width="10px", scale=0, visible=visible)
     return (
         "",
         image, 
         source_languaje, 
         target_languaje, 
-        get_audio_button, 
-        transcribe_audio_button,
+        translate_button, 
         original_audio, 
         original_audio_transcribed, 
         translated_audio, 
@@ -66,14 +68,13 @@ def get_youtube_video_id(url):
             return None
 
 def is_valid_url(url):
-    source_languaje = gr.Dropdown(visible=True, label="Source languaje", show_label=True, value="English", choices=ORIGINAL_LANGUAGE_NAME_TO_CODE, scale=1, interactive=True)
-    target_languaje = gr.Dropdown(visible=True, label="Target languaje", show_label=True, value="Español", choices=S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES, scale=1, interactive=True)
-    get_audio_button = gr.Button(size="lg", value="get audio", min_width="10px", scale=0, visible=True)
+    source_languaje = gr.Dropdown(visible=True, label="Source languaje", show_label=True, value="English", choices=language_dict, scale=1, interactive=True)
+    target_languaje = gr.Dropdown(visible=True, label="Target languaje", show_label=True, value="Español", choices=language_dict, scale=1, interactive=True)
+    translate_button = gr.Button(size="lg", value="translate", min_width="10px", scale=0, visible=True)
     original_audio = gr.Audio(label="Original audio", elem_id="original_audio", visible=True, interactive=False)
     original_audio_transcribed = gr.Textbox(label="Original audio transcribed", elem_id="original_audio_transcribed", interactive=False, visible=True)
     original_audio_translated = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", interactive=False, visible=True)
     translated_audio = gr.Audio(label="Translated audio", elem_id="translated_audio", visible=True)
-    transcribe_audio_button = gr.Button(size="lg", value="transcribe audio", min_width="10px", scale=0, visible=True)
     if "youtube" in url.lower() or "youtu.be" in url.lower():
         thumbnail = get_youtube_video_id(url)
         if thumbnail:
@@ -81,12 +82,11 @@ def is_valid_url(url):
                 gr.Image(value=thumbnail, visible=True, show_download_button=False, container=False), 
                 source_languaje,
                 target_languaje,
-                get_audio_button, 
+                translate_button, 
                 gr.Textbox(value=YOUTUBE, label="Stream page", elem_id="stream_page", visible=False),
                 original_audio,
                 original_audio_transcribed, 
                 translated_audio,
-                transcribe_audio_button, 
                 original_audio_translated, 
             )
         else:
@@ -94,12 +94,11 @@ def is_valid_url(url):
                 gr.Image(value="assets/youtube-no-thumbnails.webp", visible=True, show_download_button=False, container=False), 
                 source_languaje,
                 target_languaje,
-                get_audio_button, 
+                translate_button, 
                 gr.Textbox(value=YOUTUBE, label="Stream page", elem_id="stream_page", visible=False),
                 original_audio,
                 original_audio_transcribed, 
                 translated_audio,
-                transcribe_audio_button, 
                 original_audio_translated, 
             )
     elif "twitch" in url.lower() or "twitch.tv" in url.lower():
@@ -107,36 +106,33 @@ def is_valid_url(url):
             gr.Image(value="assets/twitch.webp", visible=True, show_download_button=False, container=False), 
             source_languaje,
             target_languaje,
-            get_audio_button, 
+            translate_button, 
             gr.Textbox(value=TWITCH, label="Stream page", elem_id="stream_page", visible=False),
             original_audio,
             original_audio_transcribed, 
             translated_audio,
-            transcribe_audio_button, 
             original_audio_translated, 
         )
     else:
         visible = False
         image = gr.Image(value="assets/youtube_error.webp", visible=visible, show_download_button=False, container=False)
-        source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=ORIGINAL_LANGUAGE_NAME_TO_CODE, scale=1, interactive=True)
-        target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES, scale=1, interactive=True)
-        get_audio_button = gr.Button(size="lg", value="get audio", min_width="10px", scale=0, visible=visible)
+        source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=language_dict, scale=1, interactive=True)
+        target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=language_dict, scale=1, interactive=True)
+        translate_button = gr.Button(size="lg", value="translate", min_width="10px", scale=0, visible=visible)
         stream_page = gr.Textbox(value=ERROR, label="Stream page", elem_id="stream_page", visible=visible)
         original_audio = gr.Audio(label="Original audio", elem_id="original_audio", visible=visible, interactive=False)
         original_audio_transcribed = gr.Textbox(label="Original audio transcribed", elem_id="original_audio_transcribed", interactive=False, visible=visible)
         original_audio_translated = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", interactive=False, visible=visible)
         translated_audio = gr.Audio(label="Translated audio", elem_id="translated_audio", visible=visible)
-        transcribe_audio_button = gr.Button(size="lg", value="transcribe audio", min_width="10px", scale=0, visible=visible)
         return (
             image, 
             source_languaje,
             target_languaje,
-            get_audio_button, 
+            translate_button, 
             stream_page,
             original_audio,
             original_audio_transcribed, 
             translated_audio,
-            transcribe_audio_button, 
             original_audio_translated, 
         )
 
@@ -175,16 +171,20 @@ def get_audio_from_video(url, stream_page):
             gr.Textbox(value=filename, label="Stream page", elem_id="stream_page", visible=False)
         )
 
-def trascribe_audio(audio_path):
+def trascribe_audio(audio_path, source_lang):
+    # Load the model
+    trascribe_model = whisper.load_model("large-v2", device=device)
+
+    # load audio and pad/trim it to fit 30 seconds
     audio = whisper.load_audio(audio_path)
     audio = whisper.pad_or_trim(audio)
     
-    mel = whisper.log_mel_spectrogram(audio).to(model.device)
-    
-    _, probs = model.detect_language(mel)
+    # make log-Mel spectrogram and move to the same device as the model
+    mel = whisper.log_mel_spectrogram(audio).to(trascribe_model.device)
     
-    options = whisper.DecodingOptions(fp16 = False)
-    result = whisper.decode(model, mel, options)
+    # Decode the result
+    options = whisper.DecodingOptions(fp16 = False, language = language_dict[source_lang]['transcriber'])
+    result = whisper.decode(trascribe_model, mel, options)
 
     # Save the result to a file
     filename = "result.txt"
@@ -192,7 +192,16 @@ def trascribe_audio(audio_path):
         f.write(result.text)
     
     # Remove audio file
-    # subprocess.run(["rm", audio_path])
+    subprocess.run(["rm", audio_path])
+
+    # free gpu memory
+    del trascribe_model
+    del audio
+    del mel
+    del options
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    gc.collect()
     
     return (
         result.text,
@@ -200,18 +209,42 @@ def trascribe_audio(audio_path):
     )
 
 def translate(original_audio_transcribed_path, source_languaje, target_languaje):
-    # Translate
+    # model
+    translate_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
+    translate_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+
+    # Get source and target languaje codes
+    source_languaje_code = language_dict[source_languaje]["translator"]
+    target_languaje_code = language_dict[target_languaje]["translator"]
+
+    # Get the transcribed text
     with open(original_audio_transcribed_path, "r") as f:
-        text = f.read()
-    translated = text
+        transcribed_text = f.read()
+    
+    # Translate the text
+    encoded = translate_tokenizer(transcribed_text, return_tensors="pt").to(device)
+    generated_tokens = translate_model.generate(
+        **encoded,
+        forced_bos_token_id=translate_tokenizer.lang_code_to_id[target_languaje_code]
+    )
+    translated = translate_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
     # Save the result to a file
     filename = "translated_text.txt"
     with open(filename, "w") as f:
-        f.write(text)
+        f.write(translated)
     
-    # Remove audio file
-    # subprocess.run(["rm", original_audio_transcribed_path])
+    # Remove transcribed file
+    subprocess.run(["rm", original_audio_transcribed_path])
+
+    # free gpu memory
+    del translate_model
+    del translate_tokenizer
+    del encoded
+    del generated_tokens
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    gc.collect()
 
     return (
         translated,
@@ -219,14 +252,26 @@ def translate(original_audio_transcribed_path, source_languaje, target_languaje)
     )
 
 def tex2speech(original_audio_translated_path):
-    pass
+    with open(original_audio_translated_path, "r") as f:
+        translated_text = f.read()
+
+    preload_models()
+    speech_array = generate_audio(translated_text, history_prompt="v2/es_speaker_1")
+
+    translated_audio_path = "translated_audio.wav"
+    write_wav(translated_audio_path, SAMPLE_RATE, speech_array)
+
+    return translated_audio_path
+
+def delete_translated_audio(translated_audio_path):
+    subprocess.run(["rm", translated_audio_path])
 
 with gr.Blocks() as demo:
+    # Layout
     with gr.Row(variant="panel"):
         url_textbox = gr.Textbox(placeholder="Add video URL here", label="Video URL", elem_id="video_url", scale=1, interactive=True)
         copy_button   = gr.Button(size="sm", icon="icons/copy.svg",   value="", min_width="10px", scale=0)
         delete_button = gr.Button(size="sm", icon="icons/delete.svg", value="", min_width="10px", scale=0)
-    copy_button.click(fn=copy_url_from_clipboard, outputs=url_textbox)
 
     stream_page = gr.Textbox(label="Stream page", elem_id="stream_page", visible=False)
     visible = False
@@ -234,11 +279,10 @@ with gr.Blocks() as demo:
         image = gr.Image(visible=visible, scale=1)
         with gr.Column():
             with gr.Row():
-                source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=ORIGINAL_LANGUAGE_NAME_TO_CODE, scale=1, interactive=True)
-                target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=S2ST_TARGET_ORIGINAL_LANGUAGE_NAMES, scale=1, interactive=True)
+                source_languaje = gr.Dropdown(visible=visible, label="Source languaje", show_label=True, value="English", choices=language_dict, scale=1, interactive=True)
+                target_languaje = gr.Dropdown(visible=visible, label="Target languaje", show_label=True, value="Español", choices=language_dict, scale=1, interactive=True)
             with gr.Row():
-                get_audio_button = gr.Button(size="lg", value="get audio", min_width="10px", scale=0, visible=visible)
-                transcribe_audio_button = gr.Button(size="lg", value="transcribe audio", min_width="10px", scale=0, visible=visible)
+                translate_button = gr.Button(size="lg", value="translate", min_width="10px", scale=0, visible=visible)
 
     original_audio = gr.Audio(label="Original audio", elem_id="original_audio", visible=visible, interactive=False)
     original_audio_path = gr.Textbox(label="Stream page", elem_id="stream_page", visible=False)
@@ -247,40 +291,44 @@ with gr.Blocks() as demo:
     original_audio_translated = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", interactive=False, visible=visible)
     original_audio_translated_path = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", visible=False)
     translated_audio = gr.Audio(label="Translated audio", elem_id="translated_audio", visible=visible)
-    url_textbox.change(
-        fn=is_valid_url, 
-        inputs=url_textbox, 
+    translated_audio_translated_path = gr.Textbox(label="translated audio translated", elem_id="translated_audio_translated", visible=False)
+
+    # Events
+    copy_button.click(fn=copy_url_from_clipboard, outputs=url_textbox)
+    delete_button.click(
+        fn=clear_video_url, 
         outputs=[
+            url_textbox, 
             image, 
             source_languaje, 
             target_languaje, 
-            get_audio_button, 
-            stream_page, 
+            translate_button, 
             original_audio, 
             original_audio_transcribed, 
             translated_audio, 
-            transcribe_audio_button, 
             original_audio_translated,
         ]
     )
-    delete_button.click(
-        fn=clear_video_url, 
+    url_textbox.change(
+        fn=is_valid_url, 
+        inputs=url_textbox, 
         outputs=[
-            url_textbox, 
             image, 
             source_languaje, 
             target_languaje, 
-            get_audio_button, 
-            transcribe_audio_button, 
+            translate_button, 
+            stream_page, 
             original_audio, 
             original_audio_transcribed, 
             translated_audio, 
             original_audio_translated,
         ]
     )
-    get_audio_button.click(fn=get_audio_from_video, inputs=[url_textbox, stream_page], outputs=[original_audio, original_audio_path])
-    original_audio.change(fn=trascribe_audio, inputs=original_audio_path, outputs=[original_audio_transcribed, original_audio_transcribed_path])
+    translate_button.click(fn=get_audio_from_video, inputs=[url_textbox, stream_page], outputs=[original_audio, original_audio_path])
+    original_audio.change(fn=trascribe_audio, inputs=[original_audio_path, source_languaje], outputs=[original_audio_transcribed, original_audio_transcribed_path])
     original_audio_transcribed.change(fn=translate, inputs=[original_audio_transcribed_path, source_languaje, target_languaje], outputs=[original_audio_translated, original_audio_translated_path])
-    original_audio_translated.change(fn=tex2speech, inputs=original_audio_translated_path, outputs=translated_audio)
+    # original_audio_translated.change(fn=tex2speech, inputs=original_audio_translated_path, outputs=translated_audio)
+    # translated_audio.change(fn=delete_translated_audio, inputs=translated_audio)
+
 
 demo.launch()