Optimized audio_processing.py with optional diarization

audio_processing.py

@@ -16,38 +16,41 @@ hf_token = os.getenv("HF_TOKEN")
 CHUNK_LENGTH = 30
 OVERLAP = 2
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Global variables for models
-device = "cuda" if torch.cuda.is_available() else "cpu"
-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, device, compute_type
-    
-    # Load Whisper model
+@spaces.GPU(duration=60)
+def load_whisper_model(model_size="small"):
+    logger.info(f"Loading Whisper model (size: {model_size})...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    compute_type = "float16" if device == "cuda" else "int8"
     try:
-        whisper_model = whisperx.load_model(model_size, device, compute_type=compute_type)
+        model = whisperx.load_model(model_size, device, compute_type=compute_type)
+        logger.info(f"Whisper model loaded successfully on {device}")
+        return model
     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)
+        model = whisperx.load_model(model_size, device, compute_type=compute_type)
+        logger.info("Whisper model loaded successfully on CPU")
+        return model
+
 
+@spaces.GPU(duration=60)
 def load_diarization_pipeline():
-    global diarization_pipeline, device
-    
-    # Try to initialize diarization pipeline
+    logger.info("Loading diarization pipeline...")
     try:
-        diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
-        if device == "cuda":
-            diarization_pipeline = diarization_pipeline.to(torch.device(device))
+        pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
+        if torch.cuda.is_available():
+            pipeline = pipeline.to(torch.device("cuda"))
+        logger.info("Diarization pipeline loaded successfully")
+        return pipeline
     except Exception as e:
         logger.warning(f"Diarization pipeline initialization failed: {str(e)}. Diarization will not be available.")
-        diarization_pipeline = None
+        return None
+
 
 def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000, overlap=OVERLAP*16000):
     chunks = []
@@ -58,18 +61,17 @@ 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:]
                 
@@ -78,11 +80,9 @@ def merge_nearby_segments(segments, time_threshold=0.5, similarity_threshold=0.7
                 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):
@@ -90,7 +90,6 @@ def get_most_common_speaker(diarization_result, start_time, end_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
@@ -104,25 +103,19 @@ def split_audio(audio, max_duration=30):
     
     return splits
 
-# Main processing function with optimizations
 @spaces.GPU(duration=60)
 def process_audio(audio_file, translate=False, model_size="small", use_diarization=True):
-    global whisper_model, diarization_pipeline
-    
-    if whisper_model is None:
-        load_models(model_size)
-    
+    logger.info(f"Starting audio processing: translate={translate}, model_size={model_size}, use_diarization={use_diarization}")
     start_time = time.time()
     
     try:
+        whisper_model = load_whisper_model(model_size)
         audio = whisperx.load_audio(audio_file)
         audio_splits = split_audio(audio)
 
-        # Perform diarization if requested and pipeline is available
         diarization_result = None
         if use_diarization:
-            if diarization_pipeline is None:
-                load_diarization_pipeline()
+            diarization_pipeline = load_diarization_pipeline()
             if diarization_pipeline is not None:
                 try:
                     diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
@@ -139,8 +132,8 @@ def process_audio(audio_file, translate=False, model_size="small", use_diarizati
             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
+                segment_start = segment["start"] + (i * 30)
+                segment_end = segment["end"] + (i * 30)
                 
                 speaker = "Unknown"
                 if diarization_result is not None: