From a7faf52ae435541c25ae26ea8ebeaec44edcd707 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Mon, 1 Jan 2024 09:53:19 +0900
Subject: [PATCH] Fix: use fp16 whisper model when bf16 unavailable

---
 transcribe.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/transcribe.py b/transcribe.py
index 275d1be..d7bb428 100644
--- a/transcribe.py
+++ b/transcribe.py
@@ -34,7 +34,12 @@ if __name__ == "__main__":
     output_file = args.output_file
     initial_prompt = args.initial_prompt
 
-    model = WhisperModel("large-v3", device="cuda", compute_type="bfloat16")
+    try:
+        model = WhisperModel("large-v3", device="cuda", compute_type="bfloat16")
+    except Exception as e:
+        # Maybe bfloat16 is not supported (e.g. in colab)
+        # I don't know actually bf16 is better than fp16 or not...
+        model = WhisperModel("large-v3", device="cuda", compute_type="float16")
 
     wav_files = [
         os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".wav")