Fixed subs and whisper device

models/subtitles.py

@@ -3,7 +3,15 @@ def save_subtitles(prediction, output_file):
     ## Save prediction to subtitles file format
     with open(output_file, "w") as f:
         for i, chunk in enumerate(prediction):
-            f.write(f"{i+1}")
-            f.write(f"{chunk['timestamp'][0]} --> {chunk['timestamp'][1]}")
+            start_time = format_time(chunk['timestamp'][0])
+            end_time = format_time(chunk['timestamp'][1])
+            f.write(f"{i+1}\n")
+            f.write(f"{start_time} --> {end_time}\n")
             f.write(f"{chunk['text']}\n")
             f.write("\n")
+
+def format_time(seconds):
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    seconds = seconds % 60
+    return f"{hours:02d}:{minutes:02d}:{seconds:04.1f}"

models/whisper.py

@@ -1,12 +1,12 @@
-import torch_directml
+# import torch_directml
 from transformers import pipeline
 
 MODEL_CHECKPOINT = "openai/whisper-small"
 CHUNK_LENGTH_S = 30
 
 def get_device():
-    # return "cpu"
-    return torch_directml.device()
+    return "cpu"
+    # return torch_directml.device()
 
 def get_pipe(device, model_checkpoint=MODEL_CHECKPOINT, chunk_length_s=CHUNK_LENGTH_S):
     return pipeline(