Update speech_processing.py

speech_processing.py

@@ -1,40 +1,43 @@
-from transformers import SeamlessM4Tv2Model, AutoProcessor
-import numpy as np
-import torch
-from pydub import AudioSegment
-
-# Load processor and model
-processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large")
-model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large")
-
-def translate_audio(audio_file):
-    if audio_file is None:
-        return "No audio file detected. Please try again."
-    
-    try:
-        # Set the device (use GPU if available)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model.to(device)
-
-        # Reset audio file pointer and load audio
-        audio = AudioSegment.from_file(audio_file, format="wav")
-        audio = audio.set_frame_rate(16000).set_channels(1)
-
-        # Convert audio to float32 NumPy array
-        audio_array = np.array(audio.get_array_of_samples()).astype(np.float32) / 32768.0
-
-        # Process input
-        audio_inputs = processor(audios=audio_array, sampling_rate=16000, return_tensors="pt")
-        audio_inputs = {key: val.to(device) for key, val in audio_inputs.items()}  # Ensure tensors are on the correct device
-
-        # Generate translation
-        output_tokens = model.generate(**audio_inputs, tgt_lang="eng", generate_speech=False)
-
-        # Extract token IDs from the generated output
-        token_ids = output_tokens.sequences
-        # Decode token IDs to text
-        translated_text_from_audio = processor.batch_decode(token_ids, skip_special_tokens=True)[0]
-
-        return translated_text_from_audio
-    except Exception as e:
-        return f"Error during audio translation: {e}"
+from transformers import SeamlessM4Tv2Model, AutoProcessor
+import numpy as np
+import torch
+from pydub import AudioSegment
+import spaces
+
+# Load processor and model
+processor = AutoProcessor.from_pretrained("facebook/seamless-m4t-v2-large")
+model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large")
+
+
+@spaces.GPU()
+def translate_audio(audio_file):
+    if audio_file is None:
+        return "No audio file detected. Please try again."
+    
+    try:
+        # Set the device (use GPU if available)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device)
+
+        # Reset audio file pointer and load audio
+        audio = AudioSegment.from_file(audio_file, format="wav")
+        audio = audio.set_frame_rate(16000).set_channels(1)
+
+        # Convert audio to float32 NumPy array
+        audio_array = np.array(audio.get_array_of_samples()).astype(np.float32) / 32768.0
+
+        # Process input
+        audio_inputs = processor(audios=audio_array, sampling_rate=16000, return_tensors="pt")
+        audio_inputs = {key: val.to(device) for key, val in audio_inputs.items()}  # Ensure tensors are on the correct device
+
+        # Generate translation
+        output_tokens = model.generate(**audio_inputs, tgt_lang="eng", generate_speech=False)
+
+        # Extract token IDs from the generated output
+        token_ids = output_tokens.sequences
+        # Decode token IDs to text
+        translated_text_from_audio = processor.batch_decode(token_ids, skip_special_tokens=True)[0]
+
+        return translated_text_from_audio
+    except Exception as e:
+        return f"Error during audio translation: {e}"