Update audio_processing.py

audio_processing.py

@@ -21,20 +21,27 @@ logger = logging.getLogger(__name__)
 
 # Global variables for models
 device = "cuda" if torch.cuda.is_available() else "cpu"
-compute_type = "float16" if device == "cuda" else "float32"
+compute_type = "float16" if device == "cuda" else "int8"
 whisper_model = None
 diarization_pipeline = None
 
 def load_models(model_size="small"):
-    global whisper_model, diarization_pipeline
+    global whisper_model, diarization_pipeline, device, compute_type
     
     # Load Whisper model
-    whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
+    try:
+        whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
+    except RuntimeError as e:
+        logger.warning(f"Failed to load Whisper model on {device}. Falling back to CPU. Error: {str(e)}")
+        device = "cpu"
+        compute_type = "int8"
+        whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
     
     # Try to initialize diarization pipeline
     try:
         diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
-        diarization_pipeline = diarization_pipeline.to(torch.device(device))
+        if device == "cuda":
+            diarization_pipeline = diarization_pipeline.to(torch.device(device))
     except Exception as e:
         logger.warning(f"Diarization pipeline initialization failed: {str(e)}. Diarization will not be available.")
         diarization_pipeline = None
@@ -136,4 +143,26 @@ def process_audio(audio_file, translate=False, model_size="small"):
         logger.error(f"An error occurred during audio processing: {str(e)}")
         raise
 
-# The merge_nearby_segments and print_results functions remain unchanged
+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