Refactor model management and transcription process: Introduced a model registry for easier management of Whisper models, added functionality to download models on startup, and streamlined the audio processing pipeline to support both chunk and segment transcriptions with improved error handling and cleanup.

app.py

@@ -112,8 +112,45 @@ def upload_data_to_r2(data, bucket_name, object_name, content_type='application/
         return False
     
 from huggingface_hub import snapshot_download
-MODEL_REPO = "deepdml/faster-whisper-large-v3-turbo-ct2"   # CT2 format
-LOCAL_DIR  = f"{CACHE_ROOT}/whisper_turbo"
+
+# -----------------------------------------------------------------------------
+# Model Management
+# -----------------------------------------------------------------------------
+MODELS = {
+    "large-v3-turbo": {
+        "repo_id": "deepdml/faster-whisper-large-v3-turbo-ct2",
+        "local_dir": f"{CACHE_ROOT}/whisper_turbo_v3"
+    },
+    "large-v3": {
+        "repo_id": "Systran/faster-whisper-large-v3",
+        "local_dir": f"{CACHE_ROOT}/whisper_large_v3"
+    },
+    "large-v2": {
+        "repo_id": "Systran/faster-whisper-large-v2",
+        "local_dir": f"{CACHE_ROOT}/whisper_large_v2"
+    },
+}
+DEFAULT_MODEL = "large-v3-turbo"
+
+def _download_model(model_name: str):
+    """Downloads a model from the hub if not already present."""
+    if model_name not in MODELS:
+        raise ValueError(f"Model '{model_name}' not found in MODELS registry.")
+    
+    model_info = MODELS[model_name]
+    if not os.path.exists(model_info["local_dir"]):
+        print(f"Downloading model '{model_name}' from {model_info['repo_id']}...")
+        snapshot_download(
+            repo_id=model_info["repo_id"],
+            local_dir=model_info["local_dir"],
+            local_dir_use_symlinks=True,
+            resume_download=True
+        )
+    return model_info["local_dir"]
+
+# Download the default model on startup
+_download_model(DEFAULT_MODEL)
+
 
 # -----------------------------------------------------------------------------
 # Audio preprocess helper (from input_and_preprocess rule)
@@ -207,6 +244,55 @@ def _write_wav(path: str, pcm: bytes, sr: int = 16000):
 
 def prepare_and_save_audio_for_model(task: dict, out_dir: str) -> dict:
     """
+    1) Decode chunk(s) to mono 16k PCM.
+    2) Run VAD to locate head/tail silence.
+    3) Trim only if head or tail >= 10s.
+    4) Save the (possibly trimmed) WAV to local file(s).
+    5) Return timing metadata, including 'trimmed_start_ms' to preserve global timestamps.
+    
+    Args:
+        task: dict containing either:
+            - "chunk": single chunk dict, or
+            - "chunk": list of chunk dicts
+        out_dir: output directory for WAV files
+    
+    Returns:
+        A wrapper dict with general fields (e.g., job_id, channel, sr, filekey)
+        and a "chunks" array containing metadata dict(s) for each processed chunk.
+        This structure is returned for both single and multiple chunk inputs.
+    """
+    chunks = task["chunk"]
+    result = {
+            "job_id": task.get("job_id", "job"),
+            "channel": task["channel"],
+            "sr": 16000,
+            "options": task.get("options", None),
+            "filekey": task.get("filekey", None),
+        }
+
+    # Handle both single chunk and multiple chunks
+    if isinstance(chunks, list):
+        # Process multiple chunks
+        results = []
+        for chunk in chunks:
+            # Create a task for each chunk
+            single_chunk_task = task.copy()
+            single_chunk_task["chunk"] = chunk
+            result = _process_single_chunk(single_chunk_task, out_dir)
+            results.append(result)
+        # Compose wrapper dict with general fields applicable to all chunks
+        result["chunks"] = results
+    else:
+        # Process single chunk and wrap in the standard response structure
+        result = _process_single_chunk(task, out_dir)
+        result["chunk"] = result
+    return result
+
+
+def _process_single_chunk(task: dict, out_dir: str) -> dict:
+    """
+    Process a single chunk - extracted from the original prepare_and_save_audio_for_model logic.
+    
     1) Decode chunk to mono 16k PCM.
     2) Run VAD to locate head/tail silence.
     3) Trim only if head or tail >= 10s.
@@ -294,17 +380,17 @@ def prepare_and_save_audio_for_model(task: dict, out_dir: str) -> dict:
     }
 
 # Download once; later runs are instant
-snapshot_download(
-    repo_id=MODEL_REPO,
-    local_dir=LOCAL_DIR,
-    local_dir_use_symlinks=True,   # saves disk space
-    resume_download=True
-)
-model_cache_path = LOCAL_DIR      # <‑‑ this is what we pass to WhisperModel
+# snapshot_download(
+#     repo_id=MODEL_REPO,
+#     local_dir=LOCAL_DIR,
+#     local_dir_use_symlinks=True,   # saves disk space
+#     resume_download=True
+# )
+# model_cache_path = LOCAL_DIR      # <‑‑ this is what we pass to WhisperModel
 
 # Lazy global holder ----------------------------------------------------------
-_whisper = None
-_batched_whisper = None
+_whisper_models = {}
+_batched_whisper_models = {}
 _diarizer = None
 _embedder = None
 
@@ -328,20 +414,32 @@ except Exception as e:
     _diarizer = None
 
 @spaces.GPU   # GPU is guaranteed to exist *inside* this function
-def _load_models():
-    global _whisper, _batched_whisper, _diarizer
-    if _whisper is None:
-        print("Loading Whisper model...")
-        _whisper = WhisperModel(
+def _load_models(model_name: str = DEFAULT_MODEL):
+    global _whisper_models, _batched_whisper_models, _diarizer
+    
+    if model_name not in _whisper_models:
+        print(f"Loading Whisper model '{model_name}'...")
+        
+        model_cache_path = _download_model(model_name)
+        
+        model = WhisperModel(
             model_cache_path,
             device="cuda",
             compute_type="float16",
         )
         
         # Create batched inference pipeline for improved performance
-        _batched_whisper = BatchedInferencePipeline(model=_whisper)
-        print("Whisper model and batched pipeline loaded successfully")
-    return _whisper, _batched_whisper, _diarizer
+        batched_model = BatchedInferencePipeline(model=model)
+        
+        _whisper_models[model_name] = model
+        _batched_whisper_models[model_name] = batched_model
+        
+        print(f"Whisper model '{model_name}' and batched pipeline loaded successfully")
+        
+    whisper = _whisper_models[model_name]
+    batched_whisper = _batched_whisper_models[model_name]
+    
+    return whisper, batched_whisper, _diarizer
 
 # -----------------------------------------------------------------------------    
 class WhisperTranscriber:
@@ -362,11 +460,11 @@ class WhisperTranscriber:
         return meta
     
     @spaces.GPU           # each call gets a GPU slice
-    def transcribe_full_audio(self, audio_path, language=None, translate=False, prompt=None, batch_size=16, base_offset_s: float = 0.0, clip_timestamps=None):
+    def transcribe_full_audio(self, audio_path, language=None, translate=False, prompt=None, batch_size=16, base_offset_s: float = 0.0, clip_timestamps=None, model_name: str = DEFAULT_MODEL, transcribe_options: dict = None):
         """Transcribe the entire audio file without speaker diarization using batched inference"""
-        whisper, batched_whisper, _ = _load_models()   # models live on the GPU
+        whisper, batched_whisper, _ = _load_models(model_name)   # models live on the GPU
         
-        print(f"Transcribing full audio with batch size {batch_size}...")
+        print(f"Transcribing full audio with '{model_name}' and batch size {batch_size}...")
         start_time = time.time()
         
         # Prepare options for batched inference
@@ -383,12 +481,14 @@ class WhisperTranscriber:
             options["vad_filter"] = False
             options["clip_timestamps"] = clip_timestamps
         else:
+            vad_options = transcribe_options.get("vad_parameters", None)
             options["vad_filter"] = True  # VAD is enabled by default for batched transcription
-            options["vad_parameters"] = VadOptions(
+            options["vad_parameters"] = VadOptions(**vad_options) if vad_options else VadOptions(
                 max_speech_duration_s=whisper.feature_extractor.chunk_length,
-                min_speech_duration_ms=150,   # ignore ultra-short blips
-                min_silence_duration_ms=150,  # split on short Mandarin pauses (if supported)                speech_pad_ms=100,
-                threshold=0.25,
+                min_speech_duration_ms=180,   # ignore ultra-short blips
+                min_silence_duration_ms=120,  # split on short Mandarin pauses (if supported)                
+                speech_pad_ms=120,
+                threshold=0.35,
                 neg_threshold=0.2,
             )
         if batch_size > 1:
@@ -439,10 +539,7 @@ class WhisperTranscriber:
         return results, detected_language
     
     # Removed audio cutting; transcription is done once on the full (preprocessed) audio
-    
-    @spaces.GPU           # each call gets a GPU slice
-    # Removed segment-wise transcription; using single full-audio transcription
-    
+        
     @spaces.GPU           # each call gets a GPU slice
     def perform_diarization(self, audio_path, num_speakers=None, base_offset_s: float = 0.0):
         """Perform speaker diarization; return segments with global timestamps and per-speaker embeddings."""
@@ -840,47 +937,215 @@ class WhisperTranscriber:
             
         return grouped_segments
     
-    @spaces.GPU           # each call gets a GPU slice
-    def process_audio_full(self, task_json, language=None, translate=False, prompt=None, group_segments=True, batch_size=16):
-        """Process a single chunk using task JSON (no diarization)."""
+    @spaces.GPU
+    def process_audio_transcribe(self, task_json, language=None, 
+                        translate=False, prompt=None, batch_size=8, model_name: str = DEFAULT_MODEL):
+        """Main processing function with diarization using task JSON for a single chunk.
+
+        Transcribes full (preprocessed) audio once, performs diarization, merges speakers into transcription.
+        """
         if not task_json or not str(task_json).strip():
             return {"error": "No JSON provided"}
         
         pre_meta = None
         try:
-            print("Starting full transcription pipeline...")
+            print("Starting new processing pipeline...")
             
             # Step 1: Preprocess per chunk JSON
             print("Preprocessing chunk JSON...")
             pre_meta = self.preprocess_from_task_json(task_json)
-            if pre_meta.get("skip"):
-                return {"segments": [], "language": "unknown", "num_speakers": 1, "transcription_method": "full_audio_batched", "batch_size": batch_size}
-            wav_path = pre_meta["out_wav_path"]
-            # Adjust timestamps by trimmed_start_ms: abs_start_ms is already global start for saved file
-            base_offset_s = float(pre_meta.get("abs_start_ms", 0)) / 1000.0
-            
-            # Step 2: Transcribe the entire audio with batching
+            transcribe_options = pre_meta.get("options", None)
+            if "chunk" in pre_meta:
+                self.transcribe_chunk(pre_meta, language, translate, prompt, batch_size, model_name, transcribe_options)
+            elif "segments" in pre_meta:
+                self.transcribe_segments(pre_meta, language, translate, prompt, batch_size, model_name, transcribe_options)
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            return {"error": f"Processing failed: {str(e)}"}
+                
+
+    @spaces.GPU
+    def transcribe_chunk(self, pre_meta, 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.
+
+        Transcribes full (preprocessed) audio once, performs diarization, merges speakers into transcription.
+        """
+        try:
+            print("Transcribing chunk...")
+            # Step 1: Preprocess per chunk JSON
+            if pre_meta["chunk"].get("skip"):
+                return {"segments": [], "language": "unknown", "num_speakers": 0, "transcription_method": "diarized_segments_batched", "batch_size": batch_size}
+            wav_path = pre_meta["chunk"]["out_wav_path"]
+            base_offset_s = float(pre_meta["chunk"].get("abs_start_ms", 0)) / 1000.0
+
+            # Step 2: Transcribe full audio once
             transcription_results, detected_language = self.transcribe_full_audio(
-                wav_path, language, translate, prompt, batch_size, base_offset_s=base_offset_s
+                wav_path, language, translate, prompt, batch_size, base_offset_s=base_offset_s, clip_timestamps=None, model_name=model_name, transcribe_options=transcribe_options
             )
             
-            # Step 3: Group segments if requested (based on time gaps and sentence endings)
-            if group_segments:
-                transcription_results = self.group_segments_by_speaker(transcription_results)
-            
-            # Step 4: Return results
-            return {
+            # Step 6: Return results
+            result = {
                 "segments": transcription_results,
                 "language": detected_language,
-                "num_speakers": 1,  # Single speaker assumption
-                "transcription_method": "full_audio_batched",
-                "batch_size": batch_size
+                "batch_size": batch_size,
             }
+            # job_id = pre_meta["job_id"]
+            # task_id = pre_meta["chunk_idx"]
+            filekey = pre_meta["filekey"]#f"ai-transcribe/split/{job_id}-{task_id}.json"
+            ret = upload_data_to_r2(json.dumps(result), "intermediate", filekey)
+            if ret:
+                return {"filekey": filekey}
+            else:
+                return {"error": "Failed to upload to R2"}
                 
         except Exception as e:
             import traceback
             traceback.print_exc()
             return {"error": f"Processing failed: {str(e)}"}
+        finally:
+            # Clean up preprocessed wav
+            if pre_meta and pre_meta["chunk"].get("out_wav_path") and os.path.exists(pre_meta["chunk"]["out_wav_path"]):
+                try:
+                    os.unlink(pre_meta["chunk"]["out_wav_path"])
+                except Exception:
+                    pass
+
+    @spaces.GPU
+    def transcribe_segments(self, pre_meta, 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.
+
+        Transcribes full (preprocessed) audio once, performs diarization, merges speakers into transcription.
+        """
+        try:
+            print("Transcribing segments...")
+            
+            # Step 1: Preprocess per chunk JSON
+            chunks = pre_meta["segments"]
+            for chunk in chunks:
+                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"]
+                base_offset_s = float(chunk.get("abs_start_ms", 0)) / 1000.0
+
+                # Step 2: Transcribe full audio once
+                transcription_results, detected_language = self.transcribe_full_audio(
+                    wav_path, language, translate, prompt, batch_size, base_offset_s=base_offset_s, clip_timestamps=None, model_name=model_name, transcribe_options=transcribe_options
+                )
+                
+                # Step 6: Return results
+                result = {
+                    "chunk_idx": chunk["chunk_idx"],
+                    "channel": chunk["channel"],
+                    "job_id": pre_meta["job_id"],
+                    "segments": transcription_results,
+                    "language": detected_language,
+                    "batch_size": batch_size,
+                }
+            # job_id = pre_meta["job_id"]
+            # task_id = pre_meta["chunk_idx"]
+            filekey = pre_meta["filekey"]#f"ai-transcribe/split/{job_id}-{task_id}.json"
+            ret = upload_data_to_r2(json.dumps(result), "intermediate", filekey)
+            if ret:
+                return {"filekey": filekey}
+            else:
+                return {"error": "Failed to upload to R2"}
+                
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            return {"error": f"Processing failed: {str(e)}"}
+        finally:
+            # Clean up preprocessed wav
+            if pre_meta and pre_meta["segments"]:
+                for chunk in pre_meta["segments"]:
+                    if chunk.get("out_wav_path") and os.path.exists(chunk["out_wav_path"]):
+                        try:
+                            os.unlink(chunk["out_wav_path"])
+                        except Exception:
+                            pass
+    
+    @spaces.GPU           # each call gets a GPU slice
+    def process_audio_diarization(self, task_json, num_speakers=0):
+        """Process audio for diarization only, returning speaker information.
+        
+        Args:
+            task_json: Task JSON containing audio processing information
+            num_speakers: Number of speakers (0 for auto-detection)
+            
+        Returns:
+            str: filekey of uploaded JSON file containing diarization results
+        """
+        if not task_json or not str(task_json).strip():
+            return {"error": "No JSON provided"}
+        
+        pre_meta = None
+        try:
+            print("Starting diarization-only pipeline...")
+            
+            # Step 1: Preprocess from task JSON
+            print("Preprocessing chunk JSON...")
+            pre_meta = self.preprocess_from_task_json(task_json)
+            if pre_meta.get("skip"):
+                # Return minimal result for skipped audio
+                task = json.loads(task_json)
+                job_id = task.get("job_id", "job")
+                task_id = str(task["chunk"]["idx"])
+                
+                result = {
+                    "num_speakers": 0,
+                    "speaker_embeddings": {}
+                }
+                
+                filekey = pre_meta["filekey"]#f"ai-transcribe/split/{job_id}-{task_id}-diarization.json"
+                ret = upload_data_to_r2(json.dumps(result), "intermediate", filekey)
+                if ret:
+                    return filekey
+                else:
+                    return {"error": "Failed to upload to R2"}
+            
+            wav_path = pre_meta["chunk"]["out_wav_path"]
+            base_offset_s = float(pre_meta["chunk"].get("abs_start_ms", 0)) / 1000.0
+            
+            # Step 2: Perform diarization
+            print("Performing diarization...")
+            start_time = time.time()
+            diarization_segments, detected_num_speakers, speaker_embeddings = self.perform_diarization(
+                wav_path, num_speakers if num_speakers > 0 else None, base_offset_s=base_offset_s
+            )
+            diarization_time = time.time() - start_time
+            print(f"Diarization completed in {diarization_time:.2f} seconds")
+            # Step 3: Compose JSON response
+            result = {
+                "num_speakers": detected_num_speakers,
+                "speaker_embeddings": speaker_embeddings,
+                "diarization_segments": diarization_segments,
+                
+            }
+            if pre_meta.get("channel", None):
+                result["channel"] = pre_meta["channel"]
+                # set channel in each diarization segment
+                for seg in diarization_segments:
+                    seg["channel"] = pre_meta["channel"]
+            
+            # Step 4: Upload to R2
+            #job_id = pre_meta["job_id"]
+            #task_id = pre_meta["chunk_idx"]
+            #filekey = f"ai-transcribe/split/{job_id}-{task_id}-diarization.json"
+            filekey = pre_meta["filekey"]
+            ret = upload_data_to_r2(json.dumps(result), "intermediate", filekey)
+            if ret:
+                # Step 5: Return filekey
+                return filekey
+            else:
+                return {"error": "Failed to upload to R2"}
+                
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            return {"error": f"Diarization processing failed: {str(e)}"}
         finally:
             # Clean up preprocessed wav
             if pre_meta and pre_meta.get("out_wav_path") and os.path.exists(pre_meta["out_wav_path"]):
@@ -888,10 +1153,10 @@ class WhisperTranscriber:
                     os.unlink(pre_meta["out_wav_path"])
                 except Exception:
                     pass
-    
+
     @spaces.GPU           # each call gets a GPU slice
     def process_audio(self, task_json, num_speakers=None, language=None, 
-                     translate=False, prompt=None, group_segments=True, batch_size=8):
+                     translate=False, prompt=None, group_segments=True, batch_size=8, model_name: str = DEFAULT_MODEL):
         """Main processing function with diarization using task JSON for a single chunk.
 
         Transcribes full (preprocessed) audio once, performs diarization, merges speakers into transcription.
@@ -929,7 +1194,7 @@ class WhisperTranscriber:
 
             # Step 2: Transcribe full audio once
             transcription_results, detected_language = self.transcribe_full_audio(
-                wav_path, language, translate, prompt, batch_size, base_offset_s=base_offset_s, clip_timestamps=None
+                wav_path, language, translate, prompt, batch_size, base_offset_s=base_offset_s, clip_timestamps=None, model_name=model_name
             )
 
             unmatched_diarization_segments = []
@@ -967,6 +1232,7 @@ class WhisperTranscriber:
                             prompt=prompt,
                             batch_size=batch_size,
                             base_offset_s=d_start,
+                            model_name=model_name
                         )
                         extra_segments.extend(seg_transcription)
                     finally:
@@ -1051,30 +1317,31 @@ def format_segments_for_display(result):
     return output
 
 @spaces.GPU
-def process_audio_gradio(task_json, num_speakers, language, translate, prompt, group_segments, use_diarization, batch_size):
+def process_audio_gradio(task_json, num_speakers, language, translate, prompt, group_segments, use_diarization, batch_size, model_name):
     """Gradio interface function"""
-    if use_diarization:
-        result = transcriber.process_audio(
+    
+    result = transcriber.process_audio_transcribe(
             task_json=task_json,
             num_speakers=num_speakers if num_speakers > 0 else None,
             language=language if language != "auto" else None,
             translate=translate,
             prompt=prompt if prompt and prompt.strip() else None,
             group_segments=group_segments,
-            batch_size=batch_size
+            batch_size=batch_size,
+            model_name=model_name
         )
-    else:
-        result = transcriber.process_audio_full(
+    '''
+        result = transcriber.process_audio_transcribe(
             task_json=task_json,
             language=language if language != "auto" else None,
             translate=translate,
             prompt=prompt if prompt and prompt.strip() else None,
-            group_segments=group_segments,
-            batch_size=batch_size
+            batch_size=batch_size,
+            model_name=model_name
         )
-    
-    formatted_output = format_segments_for_display(result)
-    return formatted_output, result
+    '''
+    #formatted_output = format_segments_for_display(result)
+    return "OK", result
 
 # Create Gradio interface
 demo = gr.Blocks(
@@ -1101,6 +1368,13 @@ with demo:
             )
             
             with gr.Accordion("⚙️ Advanced Settings", open=False):
+                model_name_dropdown = gr.Dropdown(
+                    label="Whisper Model",
+                    choices=list(MODELS.keys()),
+                    value=DEFAULT_MODEL,
+                    info="Select the Whisper model to use for transcription."
+                )
+
                 use_diarization = gr.Checkbox(
                     label="Enable Speaker Diarization",
                     value=True,
@@ -1178,7 +1452,8 @@ with demo:
             prompt,
             group_segments,
             use_diarization,
-            batch_size
+            batch_size,
+            model_name_dropdown
         ],
         outputs=[output_text, output_json]
     )