Adding fastrtc

app.py

@@ -8,7 +8,6 @@ import os
 import urllib.request
 import torchaudio
 from scipy.spatial.distance import cosine
-from RealtimeSTT import AudioToTextRecorder
 import json
 import io
 import wave
@@ -126,14 +125,13 @@ class AudioProcessor:
     def __init__(self, encoder):
         self.encoder = encoder
     
-    def extract_embedding(self, audio_int16):
+    def extract_embedding(self, audio_float):
         try:
-            float_audio = audio_int16.astype(np.float32) / 32768.0
+            # Ensure audio is in the right format
+            if np.abs(audio_float).max() > 1.0:
+                audio_float = audio_float / np.abs(audio_float).max()
             
-            if np.abs(float_audio).max() > 1.0:
-                float_audio = float_audio / np.abs(float_audio).max()
-            
-            embedding = self.encoder.embed_utterance(float_audio)
+            embedding = self.encoder.embed_utterance(audio_float)
             
             return embedding
         except Exception as e:
@@ -271,52 +269,58 @@ class SpeakerChangeDetector:
         }
 
 
-class WebRTCAudioProcessor:
-    """Processes WebRTC audio streams for speaker diarization"""
-    def __init__(self, diarization_system):
-        self.diarization_system = diarization_system
-        self.audio_buffer = []
-        self.buffer_lock = threading.Lock()
-        self.processing_thread = None
-        self.is_processing = False
+class WhisperTranscriber:
+    """Simple Whisper transcriber for audio chunks"""
+    def __init__(self, model_name="distil-large-v3"):
+        self.model = None
+        self.processor = None
+        self.model_name = model_name
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
-    def process_audio(self, audio_data, sample_rate):
-        """Process incoming audio data from WebRTC"""
+    def load_model(self):
+        """Load Whisper model"""
         try:
-            # Convert audio data to numpy array if needed
-            if isinstance(audio_data, bytes):
-                audio_array = np.frombuffer(audio_data, dtype=np.int16)
-            elif isinstance(audio_data, tuple):
-                # Handle tuple format (sample_rate, audio_array)
-                sample_rate, audio_array = audio_data
-                if isinstance(audio_array, np.ndarray):
-                    if audio_array.dtype != np.int16:
-                        audio_array = (audio_array * 32767).astype(np.int16)
-                else:
-                    audio_array = np.array(audio_array, dtype=np.int16)
-            else:
-                audio_array = np.array(audio_data, dtype=np.int16)
+            from transformers import WhisperProcessor, WhisperForConditionalGeneration
             
-            # Ensure mono audio
-            if len(audio_array.shape) > 1:
-                audio_array = audio_array[:, 0]
+            self.processor = WhisperProcessor.from_pretrained(f"distil-whisper/{self.model_name}")
+            self.model = WhisperForConditionalGeneration.from_pretrained(f"distil-whisper/{self.model_name}")
+            self.model.to(self.device)
+            
+            return True
+        except Exception as e:
+            print(f"Error loading Whisper model: {e}")
+            return False
+    
+    def transcribe(self, audio_array, sample_rate=16000):
+        """Transcribe audio array"""
+        try:
+            if self.model is None:
+                return ""
+            
+            # Ensure audio is the right sample rate
+            if sample_rate != 16000:
+                audio_array = torchaudio.functional.resample(
+                    torch.tensor(audio_array).float(),
+                    orig_freq=sample_rate,
+                    new_freq=16000
+                ).numpy()
+            
+            # Process audio
+            inputs = self.processor(audio_array, sampling_rate=16000, return_tensors="pt")
+            inputs = inputs.to(self.device)
+            
+            # Generate transcription
+            with torch.no_grad():
+                predicted_ids = self.model.generate(inputs["input_features"])
+            
+            # Decode transcription
+            transcription = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)
+            
+            return transcription[0] if transcription else ""
             
-            # Add to buffer
-            with self.buffer_lock:
-                self.audio_buffer.extend(audio_array)
-                
-                # Process buffer when it's large enough (1 second of audio)
-                if len(self.audio_buffer) >= sample_rate:
-                    buffer_to_process = np.array(self.audio_buffer[:sample_rate])
-                    self.audio_buffer = self.audio_buffer[sample_rate//2:]  # Keep 50% overlap
-                    
-                    # Feed to recorder in separate thread
-                    if self.diarization_system.recorder:
-                        audio_bytes = buffer_to_process.tobytes()
-                        self.diarization_system.recorder.feed_audio(audio_bytes)
-                        
         except Exception as e:
-            print(f"Error processing WebRTC audio: {e}")
+            print(f"Transcription error: {e}")
+            return ""
 
 
 class RealtimeSpeakerDiarization:
@@ -324,86 +328,112 @@ class RealtimeSpeakerDiarization:
         self.encoder = None
         self.audio_processor = None
         self.speaker_detector = None
-        self.recorder = None
-        self.webrtc_processor = None
+        self.transcriber = None
+        self.audio_buffer = []
+        self.processing_thread = None
         self.sentence_queue = queue.Queue()
         self.full_sentences = []
         self.sentence_speakers = []
         self.pending_sentences = []
         self.displayed_text = ""
-        self.last_realtime_text = ""
         self.is_running = False
         self.change_threshold = DEFAULT_CHANGE_THRESHOLD
         self.max_speakers = DEFAULT_MAX_SPEAKERS
+        self.audio_chunks = []
+        self.chunk_counter = 0
         
     def initialize_models(self):
-        """Initialize the speaker encoder model"""
+        """Initialize the speaker encoder and transcription models"""
         try:
             device_str = "cuda" if torch.cuda.is_available() else "cpu"
             print(f"Using device: {device_str}")
             
+            # Initialize speaker encoder
             self.encoder = SpeechBrainEncoder(device=device_str)
-            success = self.encoder.load_model()
+            encoder_success = self.encoder.load_model()
+            
+            # Initialize transcriber
+            self.transcriber = WhisperTranscriber(FINAL_TRANSCRIPTION_MODEL)
+            transcriber_success = self.transcriber.load_model()
             
-            if success:
+            if encoder_success and transcriber_success:
                 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.webrtc_processor = WebRTCAudioProcessor(self)
-                print("ECAPA-TDNN model loaded successfully!")
+                print("Models loaded successfully!")
                 return True
             else:
-                print("Failed to load ECAPA-TDNN model")
+                print("Failed to load models")
                 return False
         except Exception as e:
             print(f"Model initialization error: {e}")
             return False
     
-    def live_text_detected(self, text):
-        """Callback for real-time transcription updates"""
-        text = text.strip()
-        if text:
-            sentence_delimiters = '.?!。'
-            prob_sentence_end = (
-                len(self.last_realtime_text) > 0
-                and text[-1] in sentence_delimiters
-                and self.last_realtime_text[-1] in sentence_delimiters
-            )
-
-            self.last_realtime_text = text
-
-            if prob_sentence_end and FAST_SENTENCE_END:
-                self.recorder.stop()
-            elif prob_sentence_end:
-                self.recorder.post_speech_silence_duration = SILENCE_THRESHS[0]
+    def process_audio_stream(self, audio_data, sample_rate):
+        """Process incoming audio stream data"""
+        if not self.is_running or self.encoder is None:
+            return
+        
+        try:
+            # Convert audio data to numpy array if needed
+            if isinstance(audio_data, tuple):
+                sample_rate, audio_array = audio_data
             else:
-                self.recorder.post_speech_silence_duration = SILENCE_THRESHS[1]
+                audio_array = audio_data
+            
+            # Ensure audio is float32 and normalized
+            if audio_array.dtype != np.float32:
+                if audio_array.dtype == np.int16:
+                    audio_array = audio_array.astype(np.float32) / 32768.0
+                else:
+                    audio_array = audio_array.astype(np.float32)
+            
+            # Ensure mono audio
+            if len(audio_array.shape) > 1 and audio_array.shape[1] > 1:
+                audio_array = np.mean(audio_array, axis=1)
+            
+            # Add to buffer
+            self.audio_buffer.extend(audio_array.flatten())
+            
+            # Process when we have enough audio (about 2 seconds)
+            target_length = int(sample_rate * 2.0)
+            if len(self.audio_buffer) >= target_length:
+                self.process_audio_chunk()
+                
+        except Exception as e:
+            print(f"Error processing audio stream: {e}")
     
-    def process_final_text(self, text):
-        """Process final transcribed text with speaker embedding"""
-        text = text.strip()
-        if text:
-            try:
-                bytes_data = self.recorder.last_transcription_bytes
-                self.sentence_queue.put((text, bytes_data))
-                self.pending_sentences.append(text)
-            except Exception as e:
-                print(f"Error processing final text: {e}")
+    def process_audio_chunk(self):
+        """Process accumulated audio chunk"""
+        try:
+            if len(self.audio_buffer) < SAMPLE_RATE:  # Need at least 1 second
+                return
+            
+            # Get audio chunk
+            audio_chunk = np.array(self.audio_buffer[:int(SAMPLE_RATE * 2)])
+            self.audio_buffer = self.audio_buffer[int(SAMPLE_RATE * 1.5):]  # Keep some overlap
+            
+            # Transcribe audio
+            transcription = self.transcriber.transcribe(audio_chunk, SAMPLE_RATE)
+            
+            if transcription.strip():
+                # Extract speaker embedding
+                speaker_embedding = self.audio_processor.extract_embedding(audio_chunk)
+                
+                # Add to queue for processing
+                self.sentence_queue.put((transcription.strip(), speaker_embedding))
+                
+        except Exception as e:
+            print(f"Error processing audio chunk: {e}")
     
     def process_sentence_queue(self):
         """Process sentences in the queue for speaker detection"""
         while self.is_running:
             try:
-                text, bytes_data = self.sentence_queue.get(timeout=1)
-                
-                # Convert audio data to int16
-                audio_int16 = np.int16(bytes_data * 32767)
-                
-                # Extract speaker embedding
-                speaker_embedding = self.audio_processor.extract_embedding(audio_int16)
+                text, speaker_embedding = self.sentence_queue.get(timeout=1)
                 
                 # Store sentence and embedding
                 self.full_sentences.append((text, speaker_embedding))
@@ -416,10 +446,6 @@ class RealtimeSpeakerDiarization:
                 speaker_id, similarity = self.speaker_detector.add_embedding(speaker_embedding)
                 self.sentence_speakers.append(speaker_id)
                 
-                # Remove from pending
-                if text in self.pending_sentences:
-                    self.pending_sentences.remove(text)
-                    
             except queue.Empty:
                 continue
             except Exception as e:
@@ -431,57 +457,20 @@ class RealtimeSpeakerDiarization:
             return "Please initialize models first!"
         
         try:
-            # Setup recorder configuration for WebRTC input
-            recorder_config = {
-                'spinner': False,
-                'use_microphone': False,  # We'll feed audio manually
-                'model': FINAL_TRANSCRIPTION_MODEL,
-                'language': TRANSCRIPTION_LANGUAGE,
-                'silero_sensitivity': SILERO_SENSITIVITY,
-                'webrtc_sensitivity': WEBRTC_SENSITIVITY,
-                'post_speech_silence_duration': SILENCE_THRESHS[1],
-                'min_length_of_recording': MIN_LENGTH_OF_RECORDING,
-                'pre_recording_buffer_duration': PRE_RECORDING_BUFFER_DURATION,
-                'min_gap_between_recordings': 0,
-                'enable_realtime_transcription': True,
-                'realtime_processing_pause': 0,
-                'realtime_model_type': REALTIME_TRANSCRIPTION_MODEL,
-                'on_realtime_transcription_update': self.live_text_detected,
-                'beam_size': FINAL_BEAM_SIZE,
-                'beam_size_realtime': REALTIME_BEAM_SIZE,
-                'buffer_size': BUFFER_SIZE,
-                'sample_rate': SAMPLE_RATE,
-            }
-
-            self.recorder = AudioToTextRecorder(**recorder_config)
-            
             # Start sentence processing thread
             self.is_running = True
-            self.sentence_thread = threading.Thread(target=self.process_sentence_queue, daemon=True)
-            self.sentence_thread.start()
-            
-            # Start transcription thread
-            self.transcription_thread = threading.Thread(target=self.run_transcription, daemon=True)
-            self.transcription_thread.start()
+            self.processing_thread = threading.Thread(target=self.process_sentence_queue, daemon=True)
+            self.processing_thread.start()
             
-            return "Recording started successfully! WebRTC audio input ready."
+            return "Recording started successfully! Start speaking into your microphone."
             
         except Exception as e:
             return f"Error starting recording: {e}"
     
-    def run_transcription(self):
-        """Run the transcription loop"""
-        try:
-            while self.is_running:
-                self.recorder.text(self.process_final_text)
-        except Exception as e:
-            print(f"Transcription error: {e}")
-    
     def stop_recording(self):
         """Stop the recording process"""
         self.is_running = False
-        if self.recorder:
-            self.recorder.stop()
+        self.audio_buffer = []
         return "Recording stopped!"
     
     def clear_conversation(self):
@@ -490,7 +479,7 @@ class RealtimeSpeakerDiarization:
         self.sentence_speakers = []
         self.pending_sentences = []
         self.displayed_text = ""
-        self.last_realtime_text = ""
+        self.audio_buffer = []
         
         if self.speaker_detector:
             self.speaker_detector = SpeakerChangeDetector(
@@ -522,6 +511,7 @@ class RealtimeSpeakerDiarization:
                 sentence_text, _ = sentence
                 if i >= len(self.sentence_speakers):
                     color = "#FFFFFF"
+                    speaker_name = "Speaker ?"
                 else:
                     speaker_id = self.sentence_speakers[i]
                     color = self.speaker_detector.get_color_for_speaker(speaker_id)
@@ -530,11 +520,6 @@ class RealtimeSpeakerDiarization:
                 sentences_with_style.append(
                     f'<span style="color:{color};"><b>{speaker_name}:</b> {sentence_text}</span>')
             
-            # Add pending sentences
-            for pending_sentence in self.pending_sentences:
-                sentences_with_style.append(
-                    f'<span style="color:#60FFFF;"><b>Processing:</b> {pending_sentence}</span>')
-            
             if sentences_with_style:
                 return "<br><br>".join(sentences_with_style)
             else:
@@ -557,6 +542,7 @@ class RealtimeSpeakerDiarization:
                 f"**Last Similarity:** {status['last_similarity']:.3f}",
                 f"**Change Threshold:** {status['threshold']:.2f}",
                 f"**Total Sentences:** {len(self.full_sentences)}",
+                f"**Audio Buffer Size:** {len(self.audio_buffer)}",
                 "",
                 "**Speaker Segment Counts:**"
             ]
@@ -614,27 +600,28 @@ def get_status():
     return diarization_system.get_status_info()
 
 
-def process_audio_stream(audio):
-    """Process audio stream from WebRTC"""
-    if diarization_system.webrtc_processor and diarization_system.is_running:
-        diarization_system.webrtc_processor.process_audio(audio, SAMPLE_RATE)
-    return None
+def process_audio(audio_data):
+    """Process audio from Gradio audio input"""
+    if audio_data is not None:
+        sample_rate, audio_array = audio_data
+        diarization_system.process_audio_stream(audio_array, sample_rate)
+    return get_conversation(), get_status()
 
 
 # Create Gradio interface
 def create_interface():
-    with gr.Blocks(title="Real-time Speaker Diarization", theme=gr.themes.Monochrome()) as app:
+    with gr.Blocks(title="Real-time Speaker Diarization", theme=gr.themes.Dark()) as app:
         gr.Markdown("# 🎤 Real-time Speech Recognition with Speaker Diarization")
-        gr.Markdown("This app performs real-time speech recognition with automatic speaker identification and color-coding using WebRTC.")
+        gr.Markdown("This app performs real-time speech recognition with automatic speaker identification and color-coding using your browser's microphone.")
         
         with gr.Row():
             with gr.Column(scale=2):
-                # WebRTC Audio Input
+                # Audio input
                 audio_input = gr.Audio(
-                    sources=["microphone"],
+                    source="microphone",
+                    type="numpy",
                     streaming=True,
-                    label="🎙️ Microphone Input",
-                    type="numpy"
+                    label="🎙️ Microphone Input"
                 )
                 
                 # Main conversation display
@@ -654,7 +641,7 @@ def create_interface():
                 status_output = gr.Textbox(
                     label="System Status",
                     value="System not initialized",
-                    lines=8,
+                    lines=10,
                     interactive=False
                 )
             
@@ -681,17 +668,6 @@ def create_interface():
                 
                 update_settings_btn = gr.Button("Update Settings")
                 
-                # Instructions
-                gr.Markdown("## 📝 Instructions")
-                gr.Markdown("""
-                1. Click **Initialize System** to load models
-                2. Click **Start Recording** to begin processing
-                3. Allow microphone access when prompted
-                4. Speak into your microphone
-                5. Watch real-time transcription with speaker labels
-                6. Adjust settings as needed
-                """)
-                
                 # Speaker color legend
                 gr.Markdown("## 🎨 Speaker Colors")
                 color_info = []
@@ -699,10 +675,18 @@ def create_interface():
                     color_info.append(f'<span style="color:{color};">■</span> Speaker {i+1} ({name})')
                 
                 gr.HTML("<br>".join(color_info[:DEFAULT_MAX_SPEAKERS]))
-        
-        # Auto-refresh conversation and status
-        def refresh_display():
-            return get_conversation(), get_status()
+                
+                # Instructions
+                gr.Markdown("""
+                ## 📋 Instructions
+                1. **Initialize System** - Load AI models
+                2. **Allow microphone access** when prompted
+                3. **Start Recording** - Begin real-time processing
+                4. **Speak naturally** - The system will detect different speakers
+                5. **Stop Recording** when done
+                
+                **Note:** Processing happens in real-time with ~2 second chunks for better accuracy.
+                """)
         
         # Event handlers
         def on_initialize():
@@ -767,17 +751,19 @@ def create_interface():
             outputs=[status_output]
         )
         
-        # Connect WebRTC audio stream to processing
+        # Process streaming audio
         audio_input.stream(
-            process_audio_stream,
+            process_audio,
             inputs=[audio_input],
-            outputs=[]
+            outputs=[conversation_output, status_output],
+            time_limit=60,
+            stream_every=0.5
         )
         
-        # Auto-refresh every 2 seconds when recording
-        refresh_timer = gr.Timer(2.0)
+        # Auto-refresh every 3 seconds
+        refresh_timer = gr.Timer(3.0)
         refresh_timer.tick(
-            refresh_display,
+            lambda: (get_conversation(), get_status()),
             outputs=[conversation_output, status_output]
         )