Update modules/whisper/whisper_base.py

modules/whisper/whisper_base.py

@@ -220,6 +220,7 @@ class WhisperBase(ABC):
             model = whisper.load_model("base")
             
             files_info = {}
+            files_to_download = {}
             for file in files:
 
                 ## Detect language
@@ -240,16 +241,16 @@ class WhisperBase(ABC):
                     *whisper_params,
                 )
                 
-#                file_name, file_ext = os.path.splitext(os.path.basename(file))
-#                subtitle, file_path = self.generate_and_write_file(
-#                    file_name=file_name,
-#                    transcribed_segments=transcribed_segments,
-#                    add_timestamp=add_timestamp,
-#                    file_format=file_format,
-#                    output_dir=self.output_dir
-#                )				
-#                
-#                files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
+                file_name, file_ext = os.path.splitext(os.path.basename(file))
+                subtitle, file_path = self.generate_and_write_file(
+                    file_name=file_name,
+                    transcribed_segments=transcribed_segments,
+                    add_timestamp=add_timestamp,
+                    file_format=file_format,
+                    output_dir=self.output_dir
+                )				
+                
+                files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
 
                 ## Add output file as txt
                 file_name, file_ext = os.path.splitext(os.path.basename(file))
@@ -260,7 +261,7 @@ class WhisperBase(ABC):
                     file_format="txt",
                     output_dir=self.output_dir
                 )				
-                files_info[file_name+"_txt"] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
+                files_to_download[file_name+"_txt"] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
                 
                 ## Add output file as srt
                 file_name, file_ext = os.path.splitext(os.path.basename(file))
@@ -271,7 +272,7 @@ class WhisperBase(ABC):
                     file_format="srt",
                     output_dir=self.output_dir
                 )				
-                files_info[file_name+"_srt"] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
+                files_to_download[file_name+"_srt"] = {"subtitle": subtitle, "time_for_task": time_for_task, "path": file_path, "lang": file_language}
             
             total_result = ''
             total_info = ''
@@ -282,12 +283,12 @@ class WhisperBase(ABC):
                 #total_info += f'{info["lang"]}'
                 total_info += f"Language {info['lang']} detected"
                 # Just get first line:
-                break
+                #break
             
             #result_str = f"Processing of file '{file_name}{file_ext}' done in {self.format_time(total_time)}:\n\n{total_result}"
             total_info += f"\nTranscription process done in {self.format_time(total_time)}"
             result_str = total_result
-            result_file_path = [info['path'] for info in files_info.values()]
+            result_file_path = [info['path'] for info in files_to_download.values()]
 
             return [result_str, result_file_path, total_info]