Update whisper_cs.py (#41)

- Update whisper_cs.py (b7e436130fa58c70354776f4df83c4b894aef18b)

whisper_cs.py

@@ -178,8 +178,99 @@ def transcribe_audio(model, audio_path: str) -> Dict:
             'error': str(e),
             'success': False
         }
-    
 
+
+
+def generate(audio_path, use_v2_fast):
+    global faster_model  
+    if DEBUG_MODE: print(f"Entering generate function...")
+    if DEBUG_MODE: print(f"use_v2_fast: {use_v2_fast}")
+
+    if use_v2_fast and torch.cuda.is_available():
+        try:
+            faster_model.to("cuda")
+            print("[INFO] Moved faster_model to CUDA")
+        except Exception as e:
+            print(f"[WARNING] Could not move model to CUDA: {e}")
+
+    if use_v2_fast:
+        split_stereo_channels(audio_path)
+        left_channel_path = "temp_mono_speaker2.wav"
+        right_channel_path = "temp_mono_speaker1.wav"
+
+        left_waveform, _ = format_audio(left_channel_path)
+        right_waveform, _ = format_audio(right_channel_path)
+
+        left_waveform = left_waveform.numpy().astype("float32")
+        right_waveform = right_waveform.numpy().astype("float32")
+
+        left_result, _ = faster_model.transcribe(left_waveform, beam_size=5, task="transcribe")
+        right_result, _ = faster_model.transcribe(right_waveform, beam_size=5, task="transcribe")
+
+        left_result = list(left_result)
+        right_result = list(right_result)
+
+        def get_faster_segments(segments, speaker_label):
+            return [
+                (seg.start, seg.end, speaker_label, post_process_transcription(seg.text.strip()))
+                for seg in segments if seg.text
+            ]
+
+        left_segs = get_faster_segments(left_result, "Speaker 1")
+        right_segs = get_faster_segments(right_result, "Speaker 2")
+
+        merged_transcript = sorted(
+            left_segs + right_segs,
+            key=lambda x: float(x[0]) if x[0] is not None else float("inf")
+        )
+
+        clean_output = ""
+        for start, end, speaker, text in merged_transcript:
+            clean_output += f"[{speaker}]: {text}\n"
+
+        if DEBUG_MODE: print(f"clean_output: {clean_output}")
+
+    else:
+        model = load_whisper_model(MODEL_PATH_V2)
+        split_stereo_channels(audio_path)
+        left_channel_path = "temp_mono_speaker2.wav"
+        right_channel_path = "temp_mono_speaker1.wav"
+
+        left_waveform, _ = format_audio(left_channel_path)
+        right_waveform, _ = format_audio(right_channel_path)
+
+        left_result = transcribe_audio(model, left_waveform)
+        right_result = transcribe_audio(model, right_waveform)
+
+        def get_segments(result, speaker_label):
+            segments = result.get("segments", [])
+            if not segments:
+                return []
+            return [
+                (seg.get("start", 0.0), seg.get("end", 0.0), speaker_label,
+                 post_process_transcription(seg.get("text", "").strip()))
+                for seg in segments if seg.get("text")
+            ]
+
+        left_segs = get_segments(left_result, "Speaker 1")
+        right_segs = get_segments(right_result, "Speaker 2")
+
+        merged_transcript = sorted(
+            left_segs + right_segs,
+            key=lambda x: float(x[0]) if x[0] is not None else float("inf")
+        )
+
+        clean_output = ""
+        for start, end, speaker, text in merged_transcript:
+            clean_output += f"[{speaker}]: {text}\n"
+
+    cleanup_temp_files("temp_mono_speaker1.wav", "temp_mono_speaker2.wav")
+
+    if DEBUG_MODE: print(f"Exiting generate function...")
+    return clean_output.strip()
+
+    
+'''
 def generate(audio_path, use_v2_fast):
 
     if DEBUG_MODE: print(f"Entering generate function...")
@@ -270,4 +361,5 @@ def generate(audio_path, use_v2_fast):
 
     if DEBUG_MODE: print(f"Exiting generate function...")
         
-    return clean_output 
+    return clean_output
+'''