marquesafonso commited on
Commit
1ae9511
·
verified ·
1 Parent(s): 15ec745

fix transcriber to deal with audio or video

Browse files
Files changed (1) hide show
  1. src/transcriber.py +6 -2
src/transcriber.py CHANGED
@@ -43,12 +43,16 @@ def write_srt(segments, max_words_per_line, srt_path):
43
  file.write(result)
44
  return result, srt_path, " ".join(result_clean)
45
 
46
- def transcriber(video_input:gr.File,
47
  max_words_per_line:int,
48
  task:str,
49
  model_version:str):
 
50
  srt_filepath = os.path.normpath(f"{video_input.split('.')[0]}.srt")
51
- audio_input = convert_video_to_audio(video_input)
 
 
 
52
  model = WhisperModel(model_version, device="cpu", compute_type="int8")
53
  segments, _ = model.transcribe(
54
  audio_input,
 
43
  file.write(result)
44
  return result, srt_path, " ".join(result_clean)
45
 
46
+ def transcriber(file_input:gr.Audio | gr.Video,
47
  max_words_per_line:int,
48
  task:str,
49
  model_version:str):
50
+
51
  srt_filepath = os.path.normpath(f"{video_input.split('.')[0]}.srt")
52
+ if type(file_input) == gr.Video:
53
+ audio_input = convert_video_to_audio(file_input)
54
+ else:
55
+ audio_input = file_input
56
  model = WhisperModel(model_version, device="cpu", compute_type="int8")
57
  segments, _ = model.transcribe(
58
  audio_input,