Update modules/whisper/whisper_base.py

modules/whisper/whisper_base.py

@@ -15,7 +15,7 @@ from copy import deepcopy
 from modules.uvr.music_separator import MusicSeparator
 from modules.utils.paths import (WHISPER_MODELS_DIR, DIARIZATION_MODELS_DIR, OUTPUT_DIR, DEFAULT_PARAMETERS_CONFIG_PATH,
                                  UVR_MODELS_DIR)
-from modules.utils.subtitle_manager import get_srt, get_vtt, get_txt, get_csv, write_file, safe_filename
+from modules.utils.subtitle_manager import get_srt, get_vtt, get_txt, get_plaintext, get_csv, write_file, safe_filename
 from modules.utils.youtube_manager import get_ytdata, get_ytaudio
 from modules.utils.files_manager import get_media_files, format_gradio_files, load_yaml, save_yaml
 from modules.whisper.whisper_parameter import *
@@ -325,9 +325,16 @@ class WhisperBase(ABC):
                     else:
                         translation_note = "Already in " + target_lang
                 
-                ## Get preview as txt
+                ## Get preview
                 file_name, file_ext = os.path.splitext(os.path.basename(file))
-                subtitle = get_txt(transcribed_segments)
+                
+                if preview_format == "text":
+                    subtitle = get_txt(transcribed_segments)
+                if preview_format == "text with timestamps":
+                    subtitle = get_plaintext(transcribed_segments)
+                if preview_format == "csv":    
+                    subtitle = get_csv(transcribed_segments)
+
                 files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task, "lang": file_language, "lang_prob": file_lang_probs, "input_source_file": (file_name+file_ext), "translation": translation_note, "transcription": transcription_note}
                 
                 ## Add output file as txt
@@ -367,7 +374,9 @@ class WhisperBase(ABC):
             total_info = ''
             total_time = 0
             for file_name, info in files_info.items():
-                total_result += f'{info["subtitle"]}'
+                #total_result += f'{info["subtitle"]}'
+                total_result = {info["subtitle"]   
+                
                 total_time += info["time_for_task"]
                 total_info += f'Input file:\t\t{info["input_source_file"]}\nLanguage:\t{info["lang"]} (probability {info["lang_prob"]}%)\n'