add translation tab

app.py

@@ -3,7 +3,12 @@ import spaces
 import torch
 from loadimg import load_img
 from torchvision import transforms
-from transformers import AutoModelForImageSegmentation, pipeline
+from transformers import (
+    AutoModelForImageSegmentation,
+    pipeline,
+    MBartForConditionalGeneration,
+    MBart50TokenizerFast,
+)
 from diffusers import FluxFillPipeline
 from PIL import Image, ImageOps
 
@@ -11,9 +16,11 @@ from PIL import Image, ImageOps
 import numpy as np
 from simple_lama_inpainting import SimpleLama
 from contextlib import contextmanager
+
 # import whisperx
 import gc
 
+
 @contextmanager
 def float32_high_matmul_precision():
     torch.set_float32_matmul_precision("high")
@@ -187,7 +194,7 @@ def erase(image=None, mask=None):
 #         model = whisperx.load_model("large-v2", device, compute_type=compute_type)
 #         audio_input = whisperx.load_audio(audio)
 #         result = model.transcribe(audio_input, batch_size=batch_size)
-        
+
 #         # Clear GPU memory
 #         del model
 #         gc.collect()
@@ -205,7 +212,7 @@ def erase(image=None, mask=None):
 #         # 3. Assign speaker labels
 #         diarize_model = whisperx.DiarizationPipeline(device=device)
 #         diarize_segments = diarize_model(audio_input)
-        
+
 #         # Combine transcription with speaker diarization
 #         result = whisperx.assign_word_speakers(diarize_segments, result)
 
@@ -214,7 +221,7 @@ def erase(image=None, mask=None):
 #         for segment in result["segments"]:
 #             if not isinstance(segment, dict):
 #                 continue
-                
+
 #             speaker = f"[Speaker {segment.get('speaker', 'Unknown')}]"
 #             start_time = f"{float(segment.get('start', 0)):.2f}"
 #             end_time = f"{float(segment.get('end', 0)):.2f}"
@@ -231,6 +238,32 @@ def erase(image=None, mask=None):
 #         torch.cuda.empty_cache()
 
 
+def translate_text(text, source_lang, target_lang):
+    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    
+    # Set source language
+    tokenizer.src_lang = source_lang
+    
+    # Encode the input text
+    encoded_text = tokenizer(text, return_tensors="pt")
+    
+    # Generate translation
+    generated_tokens = model.generate(
+        **encoded_text,
+        forced_bos_token_id=tokenizer.lang_code_to_id[target_lang]
+    )
+    
+    # Decode the generated tokens
+    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    
+    # Clear GPU memory
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    return translation
+
 @spaces.GPU(duration=120)
 def main(*args):
     api_num = args[0]
@@ -247,6 +280,8 @@ def main(*args):
         return erase(*args)
     # elif api_num == 6:
     #     return transcribe(*args)
+    elif api_num == 7:
+        return translate_text(*args)
 
 
 rmbg_tab = gr.Interface(
@@ -349,7 +384,49 @@ transcribe_tab = gr.Interface(
     title="Audio Transcription",
     description="Upload an audio file to extract text using WhisperX with speaker diarization",
     api_name="transcribe",
-    examples=[]
+    examples=[],
+)
+
+translate_tab = gr.Interface(
+    fn=main,
+    inputs=[
+        gr.Number(value=7, interactive=False),
+        gr.Textbox(label="Text to translate"),
+        gr.Dropdown(
+            choices=[
+                "ar_AR", "cs_CZ", "de_DE", "en_XX", "es_XX", "et_EE", "fi_FI", "fr_XX", 
+                "gu_IN", "hi_IN", "it_IT", "ja_XX", "kk_KZ", "ko_KR", "lt_LT", "lv_LV", 
+                "my_MM", "ne_NP", "nl_XX", "ro_RO", "ru_RU", "si_LK", "tr_TR", "vi_VN", 
+                "zh_CN", "af_ZA", "az_AZ", "bn_IN", "fa_IR", "he_IL", "hr_HR", "id_ID", 
+                "ka_GE", "km_KH", "mk_MK", "ml_IN", "mn_MN", "mr_IN", "pl_PL", "ps_AF", 
+                "pt_XX", "sv_SE", "sw_KE", "ta_IN", "te_IN", "th_TH", "tl_XX", "uk_UA", 
+                "ur_PK", "xh_ZA", "gl_ES", "sl_SI"
+            ],
+            label="Source Language",
+            value="en_XX"
+        ),
+        gr.Dropdown(
+            choices=[
+                "ar_AR", "cs_CZ", "de_DE", "en_XX", "es_XX", "et_EE", "fi_FI", "fr_XX", 
+                "gu_IN", "hi_IN", "it_IT", "ja_XX", "kk_KZ", "ko_KR", "lt_LT", "lv_LV", 
+                "my_MM", "ne_NP", "nl_XX", "ro_RO", "ru_RU", "si_LK", "tr_TR", "vi_VN", 
+                "zh_CN", "af_ZA", "az_AZ", "bn_IN", "fa_IR", "he_IL", "hr_HR", "id_ID", 
+                "ka_GE", "km_KH", "mk_MK", "ml_IN", "mn_MN", "mr_IN", "pl_PL", "ps_AF", 
+                "pt_XX", "sv_SE", "sw_KE", "ta_IN", "te_IN", "th_TH", "tl_XX", "uk_UA", 
+                "ur_PK", "xh_ZA", "gl_ES", "sl_SI"
+            ],
+            label="Target Language",
+            value="fr_XX"
+        )
+    ],
+    outputs=gr.Textbox(label="Translated Text"),
+    title="Text Translation",
+    description="Translate text between multiple languages using mBART-50",
+    api_name="translate",
+    examples=[
+        [7, "Hello, how are you?", "en_XX", "fr_XX"],
+        [7, "Bonjour, comment allez-vous?", "fr_XX", "en_XX"]
+    ]
 )
 
 demo = gr.TabbedInterface(
@@ -357,20 +434,19 @@ demo = gr.TabbedInterface(
         rmbg_tab,
         outpaint_tab,
         inpaint_tab,
-        #  sam2_tab,
         erase_tab,
         transcribe_tab,
+        translate_tab
     ],
     [
         "remove background",
         "outpainting",
         "inpainting",
-        #  "sam2",
         "erase",
-        # "transcribe",
+        "transcribe",
+        "translate"
     ],
     title="Utilities that require GPU",
 )
 
-
-demo.launch()
+demo.launch()

requirements.txt

@@ -3,7 +3,7 @@ spaces
 torch
 torchvision
 git+https://github.com/huggingface/diffusers.git
-transformers
+transformers>=4.30.0
 safetensors
 accelerate
 sentencepiece
@@ -22,4 +22,5 @@ einops
 # git+https://github.com/facebookresearch/sam2.git
 matplotlib
 simple-lama-inpainting
-# git+https://github.com/m-bain/whisperX.git
+# git+https://github.com/m-bain/whisperX.git
+sacremoses