Update app.py

app.py

@@ -39,8 +39,10 @@ This Gradio demo showcases **IndicSeamlessM4T**, a fine-tuned **SeamlessM4T** mo
 """
 
 hf_token = os.getenv("HF_TOKEN")
+device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
 
-model = SeamlessM4Tv2ForSpeechToText.from_pretrained("ai4bharat/seamless-m4t-v2-large-stt", torch_dtype=torch.float16, token=hf_token).to("cuda")
+model = SeamlessM4Tv2ForSpeechToText.from_pretrained("ai4bharat/seamless-m4t-v2-large-stt", torch_dtype=torch_dtype, token=hf_token).to(device)
 processor = SeamlessM4TFeatureExtractor.from_pretrained("ai4bharat/seamless-m4t-v2-large-stt", token=hf_token)
 tokenizer = SeamlessM4TTokenizer.from_pretrained("ai4bharat/seamless-m4t-v2-large-stt", token=hf_token)
 
@@ -50,17 +52,6 @@ AUDIO_SAMPLE_RATE = 16000.0
 MAX_INPUT_AUDIO_LENGTH = 60  # in seconds
 DEFAULT_TARGET_LANGUAGE = "Hindi"
 
-if torch.cuda.is_available():
-    device = torch.device("cuda:0")
-    dtype = torch.float16
-else:
-    device = torch.device("cpu")
-    dtype = torch.float32
-
-
-
-
-
 def preprocess_audio(input_audio: str) -> None:
     arr, org_sr = torchaudio.load(input_audio)
     new_arr = torchaudio.functional.resample(arr, orig_freq=org_sr, new_freq=AUDIO_SAMPLE_RATE)
@@ -78,7 +69,7 @@ def run_s2tt(input_audio: str, source_language: str, target_language: str) -> st
 
     input_audio, orig_freq = torchaudio.load(input_audio)
     input_audio = torchaudio.functional.resample(input_audio, orig_freq=orig_freq, new_freq=16000)
-    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device="cuda",dtype=torch.float16)
+    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device="cuda", dtype=torch_dtype)
 
     text_out = model.generate(**audio_inputs, tgt_lang=target_language_code)[0].float().cpu().numpy().squeeze()
 
@@ -91,7 +82,7 @@ def run_asr(input_audio: str, target_language: str) -> str:
 
     input_audio, orig_freq = torchaudio.load(input_audio)
     input_audio = torchaudio.functional.resample(input_audio, orig_freq=orig_freq, new_freq=16000)
-    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device="cuda",dtype=torch.float16)
+    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device="cuda", dtype=torch_dtype)
 
     text_out = model.generate(**audio_inputs, tgt_lang=target_language_code)[0].float().cpu().numpy().squeeze()