Refactor audio processing: Simplified the handling of audio chunks in prepare_and_save_audio_for_model and updated preprocess_from_task_json to support both single and multiple chunk tasks, enhancing flexibility in audio preparation.

app.py

@@ -269,24 +269,8 @@ def prepare_and_save_audio_for_model(task: dict, out_dir: str) -> dict:
             "options": task.get("options", None),
             "filekey": task.get("filekey", None),
         }
-
-    # Handle both single chunk and multiple chunks
-    if task.get("segments", None):
-        # Process multiple chunks
-        chunks = task["segments"]
-        results = []
-        for chunk in chunks:
-            # Create a task for each chunk
-            single_chunk_task = task.copy()
-            single_chunk_task["chunk"] = chunk
-            chunk_result = _process_single_chunk(single_chunk_task, out_dir)
-            results.append(chunk_result)
-        # Compose wrapper dict with general fields applicable to all chunks
-        result["segments"] = results
-    else:
-        # Process single chunk and wrap in the standard response structure
-        chunk_result = _process_single_chunk(task, out_dir)
-        result["chunk"] = chunk_result
+    chunk_result = _process_single_chunk(task, out_dir)
+    result["chunk"] = chunk_result
     return result
 
 
@@ -450,7 +434,7 @@ class WhisperTranscriber:
         # do **not** create the models here!
         pass
         
-    def preprocess_from_task_json(self, task_json: str) -> dict:
+    def preprocess_from_task_json(self, task_json: str) -> any:
         """Parse task JSON and run prepare_and_save_audio_for_model, returning metadata."""
         try:
             task = json.loads(task_json)
@@ -459,7 +443,14 @@ class WhisperTranscriber:
 
         out_dir = os.path.join(CACHE_ROOT, "preprocessed")
         os.makedirs(out_dir, exist_ok=True)
-        meta = prepare_and_save_audio_for_model(task, out_dir)
+        meta = None
+        #task could be a single chunk or a list of chunks
+        if isinstance(task, list):
+            meta = []
+            for chunk in task:
+                meta.append(prepare_and_save_audio_for_model(chunk, out_dir))
+        else:
+            meta = prepare_and_save_audio_for_model(task, out_dir)
         return meta
     
     @spaces.GPU           # each call gets a GPU slice
@@ -958,10 +949,10 @@ class WhisperTranscriber:
             print("Preprocessing chunk JSON...")
             pre_meta = self.preprocess_from_task_json(task_json)
             transcribe_options = pre_meta.get("options", None)
-            if "chunk" in pre_meta:
-                return self.transcribe_chunk(pre_meta, language, translate, prompt, batch_size, model_name, transcribe_options)
-            elif "segments" in pre_meta:
+            if isinstance(pre_meta, list):
                 return self.transcribe_segments(pre_meta, language, translate, prompt, batch_size, model_name, transcribe_options)
+            elif isinstance(pre_meta, dict) and "chunk" in pre_meta:
+                return self.transcribe_chunk(pre_meta, language, translate, prompt, batch_size, model_name, transcribe_options)
         except Exception as e:
             import traceback
             traceback.print_exc()
@@ -1016,7 +1007,7 @@ class WhisperTranscriber:
                     pass
 
     @spaces.GPU
-    def transcribe_segments(self, pre_meta, language=None, 
+    def transcribe_segments(self, pre_metas, language=None, 
                         translate=False, prompt=None, batch_size=8, model_name: str = DEFAULT_MODEL, transcribe_options: dict = None):
         """Main processing function with diarization using task JSON for a single chunk.
 
@@ -1026,8 +1017,8 @@ class WhisperTranscriber:
             print("Transcribing segments...")
             transcription_results = []
             # Step 1: Preprocess per chunk JSON
-            chunks = pre_meta["segments"]
-            for chunk in chunks:
+            for pre_meta in pre_metas:
+                chunk = pre_meta["chunk"]
                 if chunk.get("skip"):
                     return {"segments": [], "language": "unknown", "num_speakers": 0, "transcription_method": "diarized_segments_batched", "batch_size": batch_size}
                 wav_path = chunk["out_wav_path"]