Update backend/main.py

backend/main.py

@@ -133,12 +133,12 @@ static_files = {
     },
 }
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large", force_download=True)
+# processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large", force_download=True)
 #cache_dir="/.cache"
 
 # PM - hardcoding temporarily as my GPU doesnt have enough vram
 # model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large").to("cpu")
-model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large", force_download=True).to(device)
+# model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large", force_download=True).to(device)
 
 
 bytes_data = bytearray()
@@ -148,6 +148,18 @@ vocoder_name = "vocoder_v2" if model_name == "seamlessM4T_v2_large" else "vocode
 clients = {}
 rooms = {}
 
+import torch
+from transformers import pipeline
+translator = pipeline("automatic-speech-recognition",
+                      "facebook/seamless-m4t-v2-large",
+                      torch_dtype=torch.float32,
+                      device="cpu")
+
+converter = pipeline("translation",
+                      "facebook/seamless-m4t-v2-large",
+                      torch_dtype=torch.float32,
+                      device="cpu")
+
 
 def get_collection_users():
     return app.database["user_records"]
@@ -297,16 +309,18 @@ async def incoming_audio(sid, data, call_id):
                 tgt_sid = next(id for id in rooms[call_id] if id != sid)
                 tgt_lang = clients[tgt_sid].target_language
                 # following example from https://github.com/facebookresearch/seamless_communication/blob/main/docs/m4t/README.md#transformers-usage
-                output_tokens = processor(audios=resampled_audio, src_lang=src_lang, return_tensors="pt")
-                model_output = model.generate(**output_tokens, tgt_lang=src_lang, generate_speech=False)[0].tolist()[0]
-                asr_text = processor.decode(model_output, skip_special_tokens=True)
+                # output_tokens = processor(audios=resampled_audio, src_lang=src_lang, return_tensors="pt")
+                # model_output = model.generate(**output_tokens, tgt_lang=src_lang, generate_speech=False)[0].tolist()[0]
+                # asr_text = processor.decode(model_output, skip_special_tokens=True)
+                asr_text = translator(resampled_audio, generate_kwargs={"tgt_lang": src_lang})['text']
                 print(f"ASR TEXT = {asr_text}")
                 # ASR TEXT => ORIGINAL TEXT
 
-                t2t_tokens = processor(text=asr_text, src_lang=src_lang, tgt_lang=tgt_lang, return_tensors="pt")
-                print(f"FIRST TYPE = {type(output_tokens)}, SECOND TYPE = {type(t2t_tokens)}")
-                translated_data = model.generate(**t2t_tokens, tgt_lang=tgt_lang, generate_speech=False)[0].tolist()[0]
-                translated_text = processor.decode(translated_data, skip_special_tokens=True)
+                # t2t_tokens = processor(text=asr_text, src_lang=src_lang, tgt_lang=tgt_lang, return_tensors="pt")
+                # print(f"FIRST TYPE = {type(output_tokens)}, SECOND TYPE = {type(t2t_tokens)}")
+                # translated_data = model.generate(**t2t_tokens, tgt_lang=tgt_lang, generate_speech=False)[0].tolist()[0]
+                # translated_text = processor.decode(translated_data, skip_special_tokens=True)
+                translated_text = converter(asr_text, src_lang=src_lang, tgt_lang=tgt_lang)
                 print(f"TRANSLATED TEXT = {translated_text}")
 
                 # BO -> send translated_text to mongodb as caption record update based on call_id