Adding the predict method from facebook/seamless_m4t

app.py

@@ -24,6 +24,54 @@ DEFAULT_TARGET_LANGUAGE = "English"
 AUDIO_SAMPLE_RATE = 16000.0
 MAX_INPUT_AUDIO_LENGTH = 60  # in seconds
 
+
+def predict(
+    task_name: str,
+    audio_source: str,
+    input_audio_mic: str | None,
+    input_audio_file: str | None,
+    input_text: str | None,
+    source_language: str | None,
+    target_language: str,
+) -> tuple[tuple[int, np.ndarray] | None, str]:
+    task_name = task_name.split()[0]
+    source_language_code = LANGUAGE_NAME_TO_CODE[source_language] if source_language else None
+    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
+
+    if task_name in ["S2ST", "S2TT", "ASR"]:
+        if audio_source == "microphone":
+            input_data = input_audio_mic
+        else:
+            input_data = input_audio_file
+
+        arr, org_sr = torchaudio.load(input_data)
+        new_arr = torchaudio.functional.resample(arr, orig_freq=org_sr, new_freq=AUDIO_SAMPLE_RATE)
+        max_length = int(MAX_INPUT_AUDIO_LENGTH * AUDIO_SAMPLE_RATE)
+        if new_arr.shape[1] > max_length:
+            new_arr = new_arr[:, :max_length]
+            gr.Warning(f"Input audio is too long. Only the first {MAX_INPUT_AUDIO_LENGTH} seconds is used.")
+
+        
+        input_data = processor(audios = new_arr, sampling_rate=AUDIO_SAMPLE_RATE, return_tensors="pt").to(device)
+    else:
+        input_data = processor(text = input_text, src_lang=source_language_code, return_tensors="pt").to(device)
+
+
+    if task_name in ["S2TT", "T2TT"]:
+        tokens_ids = model.generate(**input_data, generate_speech=False, tgt_lang=target_language_code, num_beams=5, do_sample=True)[0].cpu().squeeze().detach().tolist()
+    else:
+        output = model.generate(**input_data, return_intermediate_token_ids=True, tgt_lang=target_language_code, num_beams=5, do_sample=True, spkr_id=LANG_TO_SPKR_ID[target_language_code][0])
+        
+        waveform = output.waveform.cpu().squeeze().detach().numpy()
+        tokens_ids = output.sequences.cpu().squeeze().detach().tolist()
+
+    text_out = processor.decode(tokens_ids, skip_special_tokens=True)
+
+    if task_name in ["S2ST", "T2ST"]:
+        return (AUDIO_SAMPLE_RATE, waveform), text_out
+    else:
+        return None, text_out
+
 def process_image_with_openai(image):
     image_data = convert_image_to_required_format(image)
     openai_api_key = config('OPENAI_API_KEY')  # Make sure to have this in your .env file