app.py

import gradio as gr
import numpy as np
import torch
import torchaudio
import time
import os
import urllib.request
from scipy.spatial.distance import cosine
import threading
import queue
from collections import deque
import asyncio
from typing import Generator, Tuple, List, Optional
import whisper
from transformers import pipeline

# Configuration parameters (keeping original models)
FINAL_TRANSCRIPTION_MODEL = "distil-large-v3"
FINAL_BEAM_SIZE = 5
REALTIME_TRANSCRIPTION_MODEL = "distil-small.en"
REALTIME_BEAM_SIZE = 5
TRANSCRIPTION_LANGUAGE = "en"
SILERO_SENSITIVITY = 0.4
WEBRTC_SENSITIVITY = 3
MIN_LENGTH_OF_RECORDING = 0.7
PRE_RECORDING_BUFFER_DURATION = 0.35

# Speaker change detection parameters
DEFAULT_CHANGE_THRESHOLD = 0.7
EMBEDDING_HISTORY_SIZE = 5
MIN_SEGMENT_DURATION = 1.0
DEFAULT_MAX_SPEAKERS = 4
ABSOLUTE_MAX_SPEAKERS = 10
SAMPLE_RATE = 16000
CHUNK_DURATION = 2.0  # Process audio in 2-second chunks

# Speaker labels
SPEAKER_LABELS = [f"Speaker {i+1}" for i in range(ABSOLUTE_MAX_SPEAKERS)]

class SpeechBrainEncoder:
    """ECAPA-TDNN encoder from SpeechBrain for speaker embeddings"""
    def __init__(self, device="cpu"):
        self.device = device
        self.model = None
        self.embedding_dim = 192
        self.model_loaded = False
        self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "speechbrain")
        os.makedirs(self.cache_dir, exist_ok=True)
    
    def load_model(self):
        """Load the ECAPA-TDNN model"""
        try:
            from speechbrain.pretrained import EncoderClassifier
            
            self.model = EncoderClassifier.from_hparams(
                source="speechbrain/spkrec-ecapa-voxceleb",
                savedir=self.cache_dir,
                run_opts={"device": self.device}
            )
            
            self.model_loaded = True
            return True
        except Exception as e:
            print(f"Error loading ECAPA-TDNN model: {e}")
            return False
    
    def embed_utterance(self, audio, sr=16000):
        """Extract speaker embedding from audio"""
        if not self.model_loaded:
            raise ValueError("Model not loaded. Call load_model() first.")
        
        try:
            if isinstance(audio, np.ndarray):
                waveform = torch.tensor(audio, dtype=torch.float32).unsqueeze(0)
            else:
                waveform = audio.unsqueeze(0)
            
            if sr != 16000:
                waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)
            
            with torch.no_grad():
                embedding = self.model.encode_batch(waveform)
                
            return embedding.squeeze().cpu().numpy()
        except Exception as e:
            print(f"Error extracting embedding: {e}")
            return np.zeros(self.embedding_dim)


class SpeakerChangeDetector:
    """Speaker change detector that supports configurable number of speakers"""
    def __init__(self, embedding_dim=192, change_threshold=DEFAULT_CHANGE_THRESHOLD, max_speakers=DEFAULT_MAX_SPEAKERS):
        self.embedding_dim = embedding_dim
        self.change_threshold = change_threshold
        self.max_speakers = min(max_speakers, ABSOLUTE_MAX_SPEAKERS)
        self.current_speaker = 0
        self.previous_embeddings = []
        self.last_change_time = time.time()
        self.mean_embeddings = [None] * self.max_speakers
        self.speaker_embeddings = [[] for _ in range(self.max_speakers)]
        self.last_similarity = 0.0
        self.active_speakers = set([0])
        
    def set_max_speakers(self, max_speakers):
        """Update the maximum number of speakers"""
        new_max = min(max_speakers, ABSOLUTE_MAX_SPEAKERS)
        
        if new_max < self.max_speakers:
            for speaker_id in list(self.active_speakers):
                if speaker_id >= new_max:
                    self.active_speakers.discard(speaker_id)
            
            if self.current_speaker >= new_max:
                self.current_speaker = 0
        
        if new_max > self.max_speakers:
            self.mean_embeddings.extend([None] * (new_max - self.max_speakers))
            self.speaker_embeddings.extend([[] for _ in range(new_max - self.max_speakers)])
        else:
            self.mean_embeddings = self.mean_embeddings[:new_max]
            self.speaker_embeddings = self.speaker_embeddings[:new_max]
        
        self.max_speakers = new_max
        
    def set_change_threshold(self, threshold):
        """Update the threshold for detecting speaker changes"""
        self.change_threshold = max(0.1, min(threshold, 0.99))
        
    def add_embedding(self, embedding, timestamp=None):
        """Add a new embedding and check if there's a speaker change"""
        current_time = timestamp or time.time()
        
        if not self.previous_embeddings:
            self.previous_embeddings.append(embedding)
            self.speaker_embeddings[self.current_speaker].append(embedding)
            if self.mean_embeddings[self.current_speaker] is None:
                self.mean_embeddings[self.current_speaker] = embedding.copy()
            return self.current_speaker, 1.0
        
        current_mean = self.mean_embeddings[self.current_speaker]
        if current_mean is not None:
            similarity = 1.0 - cosine(embedding, current_mean)
        else:
            similarity = 1.0 - cosine(embedding, self.previous_embeddings[-1])
        
        self.last_similarity = similarity
        
        time_since_last_change = current_time - self.last_change_time
        is_speaker_change = False
        
        if time_since_last_change >= MIN_SEGMENT_DURATION:
            if similarity < self.change_threshold:
                best_speaker = self.current_speaker
                best_similarity = similarity
                
                for speaker_id in range(self.max_speakers):
                    if speaker_id == self.current_speaker:
                        continue
                        
                    speaker_mean = self.mean_embeddings[speaker_id]
                    
                    if speaker_mean is not None:
                        speaker_similarity = 1.0 - cosine(embedding, speaker_mean)
                        
                        if speaker_similarity > best_similarity:
                            best_similarity = speaker_similarity
                            best_speaker = speaker_id
                
                if best_speaker != self.current_speaker:
                    is_speaker_change = True
                    self.current_speaker = best_speaker
                elif len(self.active_speakers) < self.max_speakers:
                    for new_id in range(self.max_speakers):
                        if new_id not in self.active_speakers:
                            is_speaker_change = True
                            self.current_speaker = new_id
                            self.active_speakers.add(new_id)
                            break
        
        if is_speaker_change:
            self.last_change_time = current_time
        
        self.previous_embeddings.append(embedding)
        if len(self.previous_embeddings) > EMBEDDING_HISTORY_SIZE:
            self.previous_embeddings.pop(0)
        
        self.speaker_embeddings[self.current_speaker].append(embedding)
        self.active_speakers.add(self.current_speaker)
        
        if len(self.speaker_embeddings[self.current_speaker]) > 30:
            self.speaker_embeddings[self.current_speaker] = self.speaker_embeddings[self.current_speaker][-30:]
            
        if self.speaker_embeddings[self.current_speaker]:
            self.mean_embeddings[self.current_speaker] = np.mean(
                self.speaker_embeddings[self.current_speaker], axis=0
            )
        
        return self.current_speaker, similarity


class AudioProcessor:
    """Processes audio data to extract speaker embeddings"""
    def __init__(self, encoder):
        self.encoder = encoder
    
    def extract_embedding(self, audio_data):
        try:
            # Ensure audio is float32 and normalized
            if audio_data.dtype != np.float32:
                audio_data = audio_data.astype(np.float32)
            
            # Normalize if needed
            if np.abs(audio_data).max() > 1.0:
                audio_data = audio_data / np.abs(audio_data).max()
            
            # Extract embedding using the loaded encoder
            embedding = self.encoder.embed_utterance(audio_data)
            
            return embedding
        except Exception as e:
            print(f"Embedding extraction error: {e}")
            return np.zeros(self.encoder.embedding_dim)


class RealTimeSpeakerDiarization:
    """Main class for real-time speaker diarization with FastRTC"""
    def __init__(self, change_threshold=DEFAULT_CHANGE_THRESHOLD, max_speakers=DEFAULT_MAX_SPEAKERS):
        self.encoder = None
        self.audio_processor = None
        self.speaker_detector = None
        self.transcription_pipeline = None
        self.change_threshold = change_threshold
        self.max_speakers = max_speakers
        self.transcript_history = []
        self.is_initialized = False
        
        # Audio processing
        self.audio_buffer = deque(maxlen=int(SAMPLE_RATE * 10))  # 10 second buffer
        self.processing_queue = queue.Queue()
        self.last_processed_time = 0
        self.current_transcript = ""
        
    def initialize(self):
        """Initialize the speaker diarization system"""
        if self.is_initialized:
            return True
            
        try:
            device_str = "cuda" if torch.cuda.is_available() else "cpu"
            print(f"Initializing models on {device_str}...")
            
            # Initialize speaker encoder
            self.encoder = SpeechBrainEncoder(device=device_str)
            success = self.encoder.load_model()
            
            if not success:
                return False
            
            # Initialize transcription pipeline
            self.transcription_pipeline = pipeline(
                "automatic-speech-recognition",
                model=f"openai/whisper-{REALTIME_TRANSCRIPTION_MODEL}",
                device=0 if torch.cuda.is_available() else -1,
                return_timestamps=True
            )
                
            self.audio_processor = AudioProcessor(self.encoder)
            self.speaker_detector = SpeakerChangeDetector(
                embedding_dim=self.encoder.embedding_dim,
                change_threshold=self.change_threshold,
                max_speakers=self.max_speakers
            )
            
            self.is_initialized = True
            print("Speaker diarization system initialized successfully!")
            return True
            
        except Exception as e:
            print(f"Initialization error: {e}")
            return False
    
    def update_settings(self, change_threshold, max_speakers):
        """Update diarization settings"""
        self.change_threshold = change_threshold
        self.max_speakers = max_speakers
        
        if self.speaker_detector:
            self.speaker_detector.set_change_threshold(change_threshold)
            self.speaker_detector.set_max_speakers(max_speakers)
    
    def process_audio_stream(self, audio_chunk, sample_rate):
        """Process real-time audio stream from FastRTC"""
        if not self.is_initialized:
            return self.get_current_transcript(), "System not initialized"
            
        try:
            # Convert to numpy array if needed
            if hasattr(audio_chunk, 'numpy'):
                audio_data = audio_chunk.numpy()
            else:
                audio_data = np.array(audio_chunk)
            
            # Handle different audio formats
            if len(audio_data.shape) > 1:
                audio_data = audio_data.mean(axis=1)  # Convert to mono
            
            # Resample if needed
            if sample_rate != SAMPLE_RATE:
                audio_data = torchaudio.functional.resample(
                    torch.tensor(audio_data), sample_rate, SAMPLE_RATE
                ).numpy()
            
            # Add to buffer
            self.audio_buffer.extend(audio_data)
            
            # Process if we have enough audio
            current_time = time.time()
            if (current_time - self.last_processed_time) >= CHUNK_DURATION:
                self.process_buffered_audio()
                self.last_processed_time = current_time
            
            return self.get_current_transcript(), f"Processing... Buffer: {len(self.audio_buffer)} samples"
            
        except Exception as e:
            error_msg = f"Error processing audio stream: {str(e)}"
            print(error_msg)
            return self.get_current_transcript(), error_msg
    
    def process_buffered_audio(self):
        """Process buffered audio for transcription and speaker diarization"""
        if len(self.audio_buffer) < int(SAMPLE_RATE * MIN_LENGTH_OF_RECORDING):
            return
            
        try:
            # Get audio data from buffer
            audio_data = np.array(list(self.audio_buffer))
            
            # Transcribe audio
            if len(audio_data) > 0:
                result = self.transcription_pipeline(
                    audio_data,
                    return_timestamps=True,
                    generate_kwargs={"language": TRANSCRIPTION_LANGUAGE}
                )
                
                transcription = result["text"].strip()
                
                if transcription and len(transcription) > 0:
                    # Extract speaker embedding
                    embedding = self.audio_processor.extract_embedding(audio_data)
                    
                    # Detect speaker
                    speaker_id, similarity = self.speaker_detector.add_embedding(embedding)
                    
                    # Format text with speaker label
                    speaker_label = SPEAKER_LABELS[speaker_id]
                    formatted_text = f"{speaker_label}: {transcription}"
                    
                    # Add to transcript
                    self.add_to_transcript(formatted_text)
                    
                    print(f"Transcribed: {formatted_text} (Similarity: {similarity:.3f})")
            
            # Clear part of the buffer to prevent memory issues
            if len(self.audio_buffer) > SAMPLE_RATE * 5:  # Keep last 5 seconds
                self.audio_buffer = deque(list(self.audio_buffer)[-SAMPLE_RATE * 3:], maxlen=int(SAMPLE_RATE * 10))
                
        except Exception as e:
            print(f"Error in process_buffered_audio: {e}")
    
    def get_current_transcript(self):
        """Get the current transcript"""
        return "\n".join(self.transcript_history) if self.transcript_history else "Listening..."
    
    def add_to_transcript(self, formatted_text: str):
        """Add formatted text to transcript history"""
        self.transcript_history.append(formatted_text)
        
        # Keep only last 50 entries to prevent memory issues
        if len(self.transcript_history) > 50:
            self.transcript_history = self.transcript_history[-50:]
    
    def clear_transcript(self):
        """Clear transcript history and reset speaker detector"""
        self.transcript_history = []
        self.audio_buffer.clear()
        if self.speaker_detector:
            self.speaker_detector = SpeakerChangeDetector(
                embedding_dim=self.encoder.embedding_dim,
                change_threshold=self.change_threshold,
                max_speakers=self.max_speakers
            )
    
    def get_status(self):
        """Get current system status"""
        if not self.is_initialized:
            return "System not initialized"
        
        if self.speaker_detector:
            active_speakers = len(self.speaker_detector.active_speakers)
            current_speaker = self.speaker_detector.current_speaker + 1
            similarity = self.speaker_detector.last_similarity
            return f"Active: {active_speakers} speakers | Current: Speaker {current_speaker} | Similarity: {similarity:.3f}"
        
        return "Ready"


# Global instance
diarization_system = RealTimeSpeakerDiarization()


def initialize_system():
    """Initialize the diarization system"""
    success = diarization_system.initialize()
    if success:
        return "✅ Speaker diarization system initialized successfully!"
    else:
        return "❌ Failed to initialize speaker diarization system. Please check your setup."


def process_realtime_audio(audio_stream, change_threshold, max_speakers):
    """Process real-time audio stream from FastRTC"""
    if not diarization_system.is_initialized:
        return "Please initialize the system first.", "System not ready"
    
    # Update settings
    diarization_system.update_settings(change_threshold, max_speakers)
    
    if audio_stream is None:
        return diarization_system.get_current_transcript(), diarization_system.get_status()
    
    # Process the audio stream
    transcript, status = diarization_system.process_audio_stream(audio_stream, SAMPLE_RATE)
    
    return transcript, diarization_system.get_status()


def clear_conversation():
    """Clear the conversation transcript"""
    diarization_system.clear_transcript()
    return "Conversation cleared. Listening...", "Ready"


def create_gradio_interface():
    """Create and return the Gradio interface with FastRTC"""
    with gr.Blocks(title="Real-time Speaker Diarization", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🎙️ Real-time Speaker Diarization with FastRTC")
        gr.Markdown("Speak into your microphone for real-time speaker diarization and transcription.")
        
        # Initialization section
        with gr.Row():
            init_btn = gr.Button("🚀 Initialize System", variant="primary", scale=1)
            init_status = gr.Textbox(label="System Status", interactive=False, scale=2)
        
        # Settings section
        with gr.Row():
            with gr.Column():
                change_threshold = gr.Slider(
                    minimum=0.1, 
                    maximum=0.9, 
                    value=DEFAULT_CHANGE_THRESHOLD,
                    step=0.05,
                    label="Speaker Change Threshold",
                    info="Lower values = more sensitive to speaker changes"
                )
            with gr.Column():
                max_speakers = gr.Slider(
                    minimum=2,
                    maximum=ABSOLUTE_MAX_SPEAKERS,
                    value=DEFAULT_MAX_SPEAKERS,
                    step=1,
                    label="Maximum Number of Speakers",
                    info="Maximum number of speakers to detect"
                )
        
        # FastRTC Audio Input
        with gr.Row():
            with gr.Column():
                # FastRTC component for real-time audio
                audio_input = gr.FastRTC(
                    audio=True,
                    video=False,
                    label="🎤 Real-time Audio Input",
                    audio_sample_rate=SAMPLE_RATE,
                    audio_channels=1
                )
                
                clear_btn = gr.Button("🗑️ Clear Conversation", variant="stop")
            
            with gr.Column():
                current_status = gr.Textbox(
                    label="Current Status",
                    interactive=False,
                    value="Click Initialize to start"
                )
        
        # Output section
        transcript_output = gr.Textbox(
            label="🔴 Live Transcript with Speaker Labels",
            lines=15,
            max_lines=25,
            interactive=False,
            value="Click Initialize, then start speaking...",
            autoscroll=True
        )
        
        # Event handlers
        init_btn.click(
            fn=initialize_system,
            outputs=[init_status]
        )
        
        # FastRTC stream processing
        audio_input.stream(
            fn=process_realtime_audio,
            inputs=[audio_input, change_threshold, max_speakers],
            outputs=[transcript_output, current_status],
            time_limit=30  # Process in 30-second chunks
        )
        
        clear_btn.click(
            fn=clear_conversation,
            outputs=[transcript_output, current_status]
        )
        
        # Instructions
        with gr.Accordion("📋 Instructions", open=False):
            gr.Markdown("""
            ## How to Use:
            
            1. **Initialize**: Click "🚀 Initialize System" to load the AI models (this may take a moment)
            2. **Allow Microphone**: Your browser will ask for microphone permission - please allow it
            3. **Adjust Settings**: 
               - **Speaker Change Threshold**: 
                 - Lower (0.3-0.5) for speakers with different voices
                 - Higher (0.6-0.8) for speakers with similar voices
               - **Max Speakers**: Set expected number of speakers (2-10)
            4. **Start Speaking**: The system will automatically transcribe and identify speakers
            5. **View Results**: See real-time transcript with speaker labels (Speaker 1, Speaker 2, etc.)
            6. **Clear**: Use "Clear Conversation" to reset and start fresh
            
            ## Features:
            - ✅ Real-time audio processing via FastRTC
            - ✅ Automatic speech recognition with Whisper
            - ✅ Speaker diarization with ECAPA-TDNN
            - ✅ Live transcript with speaker labels
            - ✅ Configurable sensitivity settings
            - ✅ Support for up to 10 speakers
            
            ## Tips:
            - Speak clearly and allow brief pauses between speakers
            - The system learns speaker characteristics over time
            - Better results with distinct speaker voices
            - Ensure good microphone quality for best performance
            """)
    
    return demo


if __name__ == "__main__":
    # Create and launch the Gradio interface
    demo = create_gradio_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )