Update audio_processing.py

audio_processing.py

@@ -36,6 +36,9 @@ def load_models(model_size="small"):
         device = "cpu"
         compute_type = "int8"
         whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
+
+def load_diarization_pipeline():
+    global diarization_pipeline, device
     
     # Try to initialize diarization pipeline
     try:
@@ -55,8 +58,55 @@ def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000, overlap=OVERLAP*16000
         chunks.append(chunk)
     return chunks
 
+def merge_nearby_segments(segments, time_threshold=0.5, similarity_threshold=0.7):
+    merged = []
+    for segment in segments:
+        if not merged or segment['start'] - merged[-1]['end'] > time_threshold:
+            merged.append(segment)
+        else:
+            # Find the overlap
+            matcher = SequenceMatcher(None, merged[-1]['text'], segment['text'])
+            match = matcher.find_longest_match(0, len(merged[-1]['text']), 0, len(segment['text']))
+            
+            if match.size / len(segment['text']) > similarity_threshold:
+                # Merge the segments
+                merged_text = merged[-1]['text'] + segment['text'][match.b + match.size:]
+                merged_translated = merged[-1].get('translated', '') + segment.get('translated', '')[match.b + match.size:]
+                
+                merged[-1]['end'] = segment['end']
+                merged[-1]['text'] = merged_text
+                if 'translated' in segment:
+                    merged[-1]['translated'] = merged_translated
+            else:
+                # If no significant overlap, append as a new segment
+                merged.append(segment)
+    return merged
+
+# Helper function to get the most common speaker in a time range
+def get_most_common_speaker(diarization_result, start_time, end_time):
+    speakers = []
+    for turn, _, speaker in diarization_result.itertracks(yield_label=True):
+        if turn.start <= end_time and turn.end >= start_time:
+            speakers.append(speaker)
+    return max(set(speakers), key=speakers.count) if speakers else "Unknown"
+
+# Helper function to split long audio files
+def split_audio(audio, max_duration=30):
+    sample_rate = 16000
+    max_samples = max_duration * sample_rate
+    
+    if len(audio) <= max_samples:
+        return [audio]
+    
+    splits = []
+    for i in range(0, len(audio), max_samples):
+        splits.append(audio[i:i+max_samples])
+    
+    return splits
+
+# Main processing function with optimizations
 @spaces.GPU
-def process_audio(audio_file, translate=False, model_size="small"):
+def process_audio_optimized(audio_file, translate=False, model_size="small", use_diarization=True):
     global whisper_model, diarization_pipeline
     
     if whisper_model is None:
@@ -66,71 +116,55 @@ def process_audio(audio_file, translate=False, model_size="small"):
     
     try:
         audio = whisperx.load_audio(audio_file)
+        audio_splits = split_audio(audio)
 
-        # Perform diarization if pipeline is available
+        # Perform diarization if requested and pipeline is available
         diarization_result = None
-        if diarization_pipeline is not None:
-            try:
-                diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
-            except Exception as e:
-                logger.warning(f"Diarization failed: {str(e)}. Proceeding without diarization.")
-
-        chunks = preprocess_audio(audio)
+        if use_diarization:
+            if diarization_pipeline is None:
+                load_diarization_pipeline()
+            if diarization_pipeline is not None:
+                try:
+                    diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
+                except Exception as e:
+                    logger.warning(f"Diarization failed: {str(e)}. Proceeding without diarization.")
 
         language_segments = []
         final_segments = []
         
-        overlap_duration = 2  # 2 seconds overlap
-        for i, chunk in enumerate(chunks):
-            chunk_start_time = i * (CHUNK_LENGTH - overlap_duration)
-            chunk_end_time = chunk_start_time + CHUNK_LENGTH
-            logger.info(f"Processing chunk {i+1}/{len(chunks)}")
-            lang = whisper_model.detect_language(chunk)
-            result_transcribe = whisper_model.transcribe(chunk, language=lang)
-            if translate:
-                result_translate = whisper_model.transcribe(chunk, task="translate")
-            chunk_start_time = i * (CHUNK_LENGTH - overlap_duration)
-            for j, t_seg in enumerate(result_transcribe["segments"]):
-                segment_start = chunk_start_time + t_seg["start"]
-                segment_end = chunk_start_time + t_seg["end"]
-                # Skip segments in the overlapping region of the previous chunk
-                if i > 0 and segment_end <= chunk_start_time + overlap_duration:
-                    print(f"Skipping segment in overlap with previous chunk: {segment_start:.2f} - {segment_end:.2f}")
-                    continue
+        for i, audio_split in enumerate(audio_splits):
+            logger.info(f"Processing split {i+1}/{len(audio_splits)}")
             
-                # Skip segments in the overlapping region of the next chunk
-                if i < len(chunks) - 1 and segment_start >= chunk_end_time - overlap_duration:
-                    print(f"Skipping segment in overlap with next chunk: {segment_start:.2f} - {segment_end:.2f}")
-                    continue
+            result = whisper_model.transcribe(audio_split)
+            lang = result["language"]
             
+            for segment in result["segments"]:
+                segment_start = segment["start"] + (i * 30)  # Adjust start time based on split
+                segment_end = segment["end"] + (i * 30)  # Adjust end time based on split
+                
                 speaker = "Unknown"
                 if diarization_result is not None:
-                    speakers = []
-                    for turn, track, spk in diarization_result.itertracks(yield_label=True):
-                        if turn.start <= segment_end and turn.end >= segment_start:
-                            speakers.append(spk)
-                    speaker = max(set(speakers), key=speakers.count) if speakers else "Unknown"
-            
-                segment = {
+                    speaker = get_most_common_speaker(diarization_result, segment_start, segment_end)
+                
+                final_segment = {
                     "start": segment_start,
                     "end": segment_end,
                     "language": lang,
                     "speaker": speaker,
-                    "text": t_seg["text"],
+                    "text": segment["text"],
                 }
                 
                 if translate:
-                    segment["translated"] = result_translate["segments"][j]["text"]
+                    translation = whisper_model.transcribe(audio_split[int(segment["start"]*16000):int(segment["end"]*16000)], task="translate")
+                    final_segment["translated"] = translation["text"]
+                
+                final_segments.append(final_segment)
             
-                final_segments.append(segment)
-
             language_segments.append({
                 "language": lang,
-                "start": chunk_start_time,
-                "end": chunk_start_time + CHUNK_LENGTH
+                "start": i * 30,
+                "end": min((i + 1) * 30, len(audio) / 16000)
             })
-            chunk_end_time = time.time()
-            logger.info(f"Chunk {i+1} processed in {chunk_end_time - chunk_start_time:.2f} seconds")
 
         final_segments.sort(key=lambda x: x["start"])
         merged_segments = merge_nearby_segments(final_segments)
@@ -143,26 +177,6 @@ def process_audio(audio_file, translate=False, model_size="small"):
         logger.error(f"An error occurred during audio processing: {str(e)}")
         raise
 
-def merge_nearby_segments(segments, time_threshold=0.5, similarity_threshold=0.7):
-    merged = []
-    for segment in segments:
-        if not merged or segment['start'] - merged[-1]['end'] > time_threshold:
-            merged.append(segment)
-        else:
-            # Find the overlap
-            matcher = SequenceMatcher(None, merged[-1]['text'], segment['text'])
-            match = matcher.find_longest_match(0, len(merged[-1]['text']), 0, len(segment['text']))
-            
-            if match.size / len(segment['text']) > similarity_threshold:
-                # Merge the segments
-                merged_text = merged[-1]['text'] + segment['text'][match.b + match.size:]
-                merged_translated = merged[-1].get('translated', '') + segment.get('translated', '')[match.b + match.size:]
-                
-                merged[-1]['end'] = segment['end']
-                merged[-1]['text'] = merged_text
-                if 'translated' in segment:
-                    merged[-1]['translated'] = merged_translated
-            else:
-                # If no significant overlap, append as a new segment
-                merged.append(segment)
-    return merged
+# You can keep the original process_audio function for backwards compatibility
+# or replace it with the optimized version
+process_audio = process_audio_optimized