Revert portg

app.py

@@ -562,53 +562,89 @@ class RealtimeSpeakerDiarization:
 
 # FastRTC Audio Handler
 # FastRTC Audio Handler for Real-time Diarization
+# FastRTC Audio Handler for Real-time Diarization
+import asyncio
+import numpy as np
+from fastrtc import AsyncStreamHandler, Stream
+from fastapi import FastAPI, APIRouter
+import gradio as gr
+import time
+import os
+import threading
+from queue import Queue
+import json
 
-
-class DiarizationHandler:
+class DiarizationHandler(AsyncStreamHandler):
     def __init__(self, diarization_system):
+        super().__init__()
         self.diarization_system = diarization_system
         self.audio_queue = Queue()
         self.is_processing = False
+        self.sample_rate = 16000  # Default sample rate
         
     def copy(self):
-        # Return a fresh handler for each new stream connection
+        """Return a fresh handler for each new stream connection"""
         return DiarizationHandler(self.diarization_system)
     
-    async def on_audio_frame(self, frame: AudioFrame):
-        """Handle incoming audio frames from FastRTC"""
+    async def emit(self):
+        """Not used in this implementation - we only receive audio"""
+        return None
+    
+    async def receive(self, frame):
+        """Receive audio data from FastRTC and process it"""
         try:
-            if self.diarization_system.is_running and frame.data is not None:
-                # Convert audio frame to numpy array
-                if isinstance(frame.data, bytes):
-                    # Convert bytes to numpy array (assuming 16-bit PCM)
-                    audio_data = np.frombuffer(frame.data, dtype=np.int16)
-                elif hasattr(frame, 'to_ndarray'):
-                    audio_data = frame.to_ndarray()
-                else:
-                    audio_data = np.array(frame.data, dtype=np.float32)
-                
-                # Ensure audio is in the right format (mono, float32, -1 to 1 range)
-                if audio_data.dtype == np.int16:
-                    audio_data = audio_data.astype(np.float32) / 32768.0
+            if not self.diarization_system.is_running:
+                return
                 
-                # If stereo, convert to mono
-                if len(audio_data.shape) > 1:
-                    audio_data = np.mean(audio_data, axis=1)
-                
-                # Feed to diarization system
-                await self.process_audio_async(audio_data, frame.sample_rate)
+            # Extract audio data from frame
+            if hasattr(frame, 'data') and frame.data is not None:
+                audio_data = frame.data
+            elif hasattr(frame, 'audio') and frame.audio is not None:
+                audio_data = frame.audio
+            else:
+                audio_data = frame
+            
+            # Convert to numpy array if needed
+            if isinstance(audio_data, bytes):
+                # Convert bytes to numpy array (assuming 16-bit PCM)
+                audio_array = np.frombuffer(audio_data, dtype=np.int16)
+                # Normalize to float32 range [-1, 1]
+                audio_array = audio_array.astype(np.float32) / 32768.0
+            elif isinstance(audio_data, (list, tuple)):
+                audio_array = np.array(audio_data, dtype=np.float32)
+            elif isinstance(audio_data, np.ndarray):
+                audio_array = audio_data.astype(np.float32)
+            else:
+                print(f"Unknown audio data type: {type(audio_data)}")
+                return
+            
+            # Ensure mono audio
+            if len(audio_array.shape) > 1 and audio_array.shape[1] > 1:
+                audio_array = np.mean(audio_array, axis=1)
+            
+            # Ensure 1D array
+            if len(audio_array.shape) > 1:
+                audio_array = audio_array.flatten()
+            
+            # Get sample rate from frame if available
+            sample_rate = getattr(frame, 'sample_rate', self.sample_rate)
+            
+            # Process audio asynchronously to avoid blocking
+            await self.process_audio_async(audio_array, sample_rate)
                 
         except Exception as e:
-            print(f"Error processing audio frame: {e}")
+            print(f"Error in FastRTC audio receive: {e}")
+            import traceback
+            traceback.print_exc()
     
     async def process_audio_async(self, audio_data, sample_rate=16000):
         """Process audio data asynchronously"""
         try:
-            # Run in thread pool to avoid blocking
+            # Run the audio processing in a thread pool to avoid blocking
             loop = asyncio.get_event_loop()
             await loop.run_in_executor(
                 None, 
-                self.diarization_system.feed_audio_data, 
+                self.diarization_system.process_audio_chunk, 
                 audio_data, 
                 sample_rate
             )
@@ -616,15 +652,19 @@ class DiarizationHandler:
             print(f"Error in async audio processing: {e}")
 
 
-# Global instance
-diarization_system = RealtimeSpeakerDiarization()
+# Global instances
+diarization_system = None  # Will be initialized when RealtimeSpeakerDiarization is available
 audio_handler = None
 
 
 def initialize_system():
     """Initialize the diarization system"""
-    global audio_handler
+    global audio_handler, diarization_system
     try:
+        if diarization_system is None:
+            print("Error: RealtimeSpeakerDiarization not initialized")
+            return "❌ Diarization system not available. Please ensure RealtimeSpeakerDiarization is properly imported."
+        
         success = diarization_system.initialize_models()
         if success:
             audio_handler = DiarizationHandler(diarization_system)
@@ -632,12 +672,15 @@ def initialize_system():
         else:
             return "❌ Failed to initialize system. Please check the logs."
     except Exception as e:
+        print(f"Initialization error: {e}")
         return f"❌ Initialization error: {str(e)}"
 
 
 def start_recording():
     """Start recording and transcription"""
     try:
+        if diarization_system is None:
+            return "❌ System not initialized"
         result = diarization_system.start_recording()
         return f"🎙️ {result} - FastRTC audio streaming is active."
     except Exception as e:
@@ -647,6 +690,8 @@ def start_recording():
 def stop_recording():
     """Stop recording and transcription"""
     try:
+        if diarization_system is None:
+            return "❌ System not initialized"
         result = diarization_system.stop_recording()
         return f"⏹️ {result}"
     except Exception as e:
@@ -656,6 +701,8 @@ def stop_recording():
 def clear_conversation():
     """Clear the conversation"""
     try:
+        if diarization_system is None:
+            return "❌ System not initialized"
         result = diarization_system.clear_conversation()
         return f"🗑️ {result}"
     except Exception as e:
@@ -665,6 +712,8 @@ def clear_conversation():
 def update_settings(threshold, max_speakers):
     """Update system settings"""
     try:
+        if diarization_system is None:
+            return "❌ System not initialized"
         result = diarization_system.update_settings(threshold, max_speakers)
         return f"⚙️ {result}"
     except Exception as e:
@@ -674,6 +723,8 @@ def update_settings(threshold, max_speakers):
 def get_conversation():
     """Get the current conversation"""
     try:
+        if diarization_system is None:
+            return "<i>System not initialized</i>"
         return diarization_system.get_formatted_conversation()
     except Exception as e:
         return f"<i>Error getting conversation: {str(e)}</i>"
@@ -682,6 +733,8 @@ def get_conversation():
 def get_status():
     """Get system status"""
     try:
+        if diarization_system is None:
+            return "System not initialized"
         return diarization_system.get_status_info()
     except Exception as e:
         return f"Error getting status: {str(e)}"
@@ -709,12 +762,17 @@ def create_interface():
                     stop_btn = gr.Button("⏹️ Stop Recording", variant="stop", size="lg", interactive=False)
                     clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="lg", interactive=False)
                 
-                # Audio connection status
+                # FastRTC Stream Interface
                 with gr.Row():
-                    connection_status = gr.HTML(
-                        value="<div style='padding: 10px; background: #fff3cd; border-radius: 5px;'>🔌 FastRTC: Not connected</div>",
-                        label="Connection Status"
-                    )
+                    gr.HTML("""
+                    <div id="fastrtc-container" style="border: 2px solid #ddd; border-radius: 10px; padding: 20px; margin: 10px 0;">
+                        <h3>🎵 Audio Stream</h3>
+                        <p>FastRTC audio stream will appear here when recording starts.</p>
+                        <div id="stream-status" style="padding: 10px; background: #f8f9fa; border-radius: 5px; margin-top: 10px;">
+                            Status: Waiting for initialization...
+                        </div>
+                    </div>
+                    """)
                 
                 # Status display
                 status_output = gr.Textbox(
@@ -749,26 +807,39 @@ def create_interface():
                 update_settings_btn = gr.Button("Update Settings", variant="secondary")
                 
                 # Audio settings
-                gr.Markdown("## 🔊 Audio Settings")
-                gr.Markdown("""
-                **Recommended settings:**
-                - Use a good quality microphone
-                - Ensure stable internet connection
-                - Speak clearly and avoid background noise
-                - Position microphone 6-12 inches from mouth
-                """)
+                gr.Markdown("## 🔊 Audio Configuration")
+                with gr.Accordion("Advanced Audio Settings", open=False):
+                    gr.Markdown("""
+                    **Current Configuration:**
+                    - Sample Rate: 16kHz
+                    - Audio Format: 16-bit PCM → Float32 (via AudioProcessor)
+                    - Channels: Mono (stereo converted automatically)
+                    - Buffer Size: 1024 samples for real-time processing
+                    - Processing: Uses existing AudioProcessor.extract_embedding()
+                    """)
                 
                 # Instructions
                 gr.Markdown("## 📝 How to Use")
                 gr.Markdown("""
                 1. **Initialize**: Click "Initialize System" to load AI models
-                2. **Connect**: Allow microphone access when prompted
-                3. **Start**: Click "Start Recording" to begin processing
-                4. **Speak**: Talk into your microphone naturally
-                5. **Monitor**: Watch real-time transcription with speaker labels
-                6. **Adjust**: Fine-tune settings as needed
+                2. **Start**: Click "Start Recording" to begin processing  
+                3. **Connect**: The FastRTC stream will activate automatically
+                4. **Allow Access**: Grant microphone permissions when prompted
+                5. **Speak**: Talk naturally into your microphone
+                6. **Monitor**: Watch real-time transcription with speaker colors
                 """)
                 
+                # Performance tips
+                with gr.Accordion("💡 Performance Tips", open=False):
+                    gr.Markdown("""
+                    - Use Chrome/Edge for best FastRTC performance
+                    - Ensure stable internet connection
+                    - Use headphones to prevent echo
+                    - Position microphone 6-12 inches away
+                    - Minimize background noise
+                    - Allow browser microphone access
+                    """)
+                
                 # Speaker color legend
                 gr.Markdown("## 🎨 Speaker Colors")
                 speaker_colors = [
@@ -784,28 +855,19 @@ def create_interface():
                 
                 color_html = ""
                 for i, (color, name) in enumerate(speaker_colors[:4]):
-                    color_html += f'<div style="display: inline-block; margin: 5px;"><span style="color:{color}; font-size: 20px;">●</span> Speaker {i+1} ({name})</div><br>'
+                    color_html += f'<div style="margin: 3px 0;"><span style="color:{color}; font-size: 16px; font-weight: bold;">●</span> Speaker {i+1} ({name})</div>'
                 
-                gr.HTML(color_html)
+                gr.HTML(f"<div style='font-size: 14px;'>{color_html}</div>")
         
         # Auto-refresh conversation and status
         def refresh_display():
             try:
                 conversation = get_conversation()
                 status = get_status()
-                
-                # Update connection status based on system state
-                if diarization_system.is_running:
-                    conn_status = "<div style='padding: 10px; background: #d4edda; border-radius: 5px;'>🟢 FastRTC: Connected & Recording</div>"
-                elif hasattr(diarization_system, 'encoder') and diarization_system.encoder is not None:
-                    conn_status = "<div style='padding: 10px; background: #d1ecf1; border-radius: 5px;'>🔵 FastRTC: Ready to connect</div>"
-                else:
-                    conn_status = "<div style='padding: 10px; background: #f8d7da; border-radius: 5px;'>🔴 FastRTC: System not initialized</div>"
-                
-                return conversation, status, conn_status
+                return conversation, status
             except Exception as e:
                 error_msg = f"Error refreshing display: {str(e)}"
-                return f"<i>{error_msg}</i>", error_msg, "<div style='padding: 10px; background: #f8d7da; border-radius: 5px;'>❌ FastRTC: Error</div>"
+                return f"<i>{error_msg}</i>", error_msg
         
         # Event handlers
         def on_initialize():
@@ -813,14 +875,13 @@ def create_interface():
                 result = initialize_system()
                 success = "successfully" in result.lower()
                 
-                conversation, status, conn_status = refresh_display()
+                conversation, status = refresh_display()
                 
                 return (
                     result,  # status_output
                     gr.update(interactive=success),   # start_btn
                     gr.update(interactive=success),   # clear_btn
                     conversation,  # conversation_output
-                    conn_status   # connection_status
                 )
             except Exception as e:
                 error_msg = f"❌ Initialization failed: {str(e)}"
@@ -829,19 +890,15 @@ def create_interface():
                     gr.update(interactive=False),
                     gr.update(interactive=False),
                     "<i>System not ready</i>",
-                    "<div style='padding: 10px; background: #f8d7da; border-radius: 5px;'>❌ FastRTC: Initialization failed</div>"
                 )
         
         def on_start():
             try:
                 result = start_recording()
-                conversation, status, conn_status = refresh_display()
-                
                 return (
                     result,  # status_output
                     gr.update(interactive=False),  # start_btn
                     gr.update(interactive=True),   # stop_btn
-                    conn_status  # connection_status
                 )
             except Exception as e:
                 error_msg = f"❌ Failed to start: {str(e)}"
@@ -849,19 +906,15 @@ def create_interface():
                     error_msg,
                     gr.update(interactive=True),
                     gr.update(interactive=False),
-                    "<div style='padding: 10px; background: #f8d7da; border-radius: 5px;'>❌ FastRTC: Start failed</div>"
                 )
         
         def on_stop():
             try:
                 result = stop_recording()
-                conversation, status, conn_status = refresh_display()
-                
                 return (
                     result,  # status_output
                     gr.update(interactive=True),   # start_btn
                     gr.update(interactive=False),  # stop_btn
-                    conn_status  # connection_status
                 )
             except Exception as e:
                 error_msg = f"❌ Failed to stop: {str(e)}"
@@ -869,13 +922,12 @@ def create_interface():
                     error_msg,
                     gr.update(interactive=False),
                     gr.update(interactive=True),
-                    "<div style='padding: 10px; background: #f8d7da; border-radius: 5px;'>❌ FastRTC: Stop failed</div>"
                 )
         
         def on_clear():
             try:
                 result = clear_conversation()
-                conversation, status, conn_status = refresh_display()
+                conversation, status = refresh_display()
                 return result, conversation
             except Exception as e:
                 error_msg = f"❌ Failed to clear: {str(e)}"
@@ -891,17 +943,17 @@ def create_interface():
         # Connect event handlers
         init_btn.click(
             on_initialize,
-            outputs=[status_output, start_btn, clear_btn, conversation_output, connection_status]
+            outputs=[status_output, start_btn, clear_btn, conversation_output]
         )
         
         start_btn.click(
             on_start,
-            outputs=[status_output, start_btn, stop_btn, connection_status]
+            outputs=[status_output, start_btn, stop_btn]
         )
         
         stop_btn.click(
             on_stop,
-            outputs=[status_output, start_btn, stop_btn, connection_status]
+            outputs=[status_output, start_btn, stop_btn]
         )
         
         clear_btn.click(
@@ -919,15 +971,15 @@ def create_interface():
         refresh_timer = gr.Timer(2.0)
         refresh_timer.tick(
             refresh_display,
-            outputs=[conversation_output, status_output, connection_status]
+            outputs=[conversation_output, status_output]
         )
     
     return interface
 
 
-# FastAPI setup for HuggingFace Spaces
+# FastAPI setup for API endpoints
 def create_fastapi_app():
-    """Create FastAPI app with proper FastRTC integration"""
+    """Create FastAPI app with API endpoints"""
     app = FastAPI(
         title="Real-time Speaker Diarization",
         description="Real-time speech recognition with speaker diarization using FastRTC",
@@ -943,8 +995,8 @@ def create_fastapi_app():
         return {
             "status": "healthy",
             "timestamp": time.time(),
-            "system_initialized": hasattr(diarization_system, 'encoder') and diarization_system.encoder is not None,
-            "recording_active": diarization_system.is_running if hasattr(diarization_system, 'is_running') else False
+            "system_initialized": diarization_system is not None and hasattr(diarization_system, 'encoder') and diarization_system.encoder is not None,
+            "recording_active": diarization_system.is_running if diarization_system and hasattr(diarization_system, 'is_running') else False
         }
     
     @router.get("/api/conversation")
@@ -954,7 +1006,7 @@ def create_fastapi_app():
             return {
                 "conversation": get_conversation(),
                 "status": get_status(),
-                "is_recording": diarization_system.is_running if hasattr(diarization_system, 'is_running') else False,
+                "is_recording": diarization_system.is_running if diarization_system and hasattr(diarization_system, 'is_running') else False,
                 "timestamp": time.time()
             }
         except Exception as e:
@@ -977,40 +1029,62 @@ def create_fastapi_app():
             
             return {
                 "result": result, 
-                "is_recording": diarization_system.is_running if hasattr(diarization_system, 'is_running') else False,
+                "is_recording": diarization_system.is_running if diarization_system and hasattr(diarization_system, 'is_running') else False,
                 "timestamp": time.time()
             }
         except Exception as e:
             return {"error": str(e), "timestamp": time.time()}
     
-    # FastRTC WebSocket endpoint for audio streaming
-    @router.websocket("/ws/audio")
-    async def websocket_audio_endpoint(websocket):
-        """WebSocket endpoint for FastRTC audio streaming"""
-        await websocket.accept()
-        
-        try:
-            while True:
-                # Receive audio data from FastRTC client
-                data = await websocket.receive_bytes()
-                
-                if audio_handler and diarization_system.is_running:
-                    # Create audio frame and process
-                    frame = AudioFrame(data=data, sample_rate=16000)
-                    await audio_handler.on_audio_frame(frame)
-                    
-        except Exception as e:
-            print(f"WebSocket error: {e}")
-        finally:
-            await websocket.close()
-    
     app.include_router(router)
     return app
 
 
-# Main application entry point
-def create_app():
-    """Create the complete application for HuggingFace Spaces"""
+# Function to setup FastRTC stream
+def setup_fastrtc_stream(app):
+    """Setup FastRTC stream with proper configuration"""
+    try:
+        if audio_handler is None:
+            print("Warning: Audio handler not initialized. Initialize system first.")
+            return None
+        
+        # Get HuggingFace token for TURN server (optional)
+        hf_token = os.environ.get("HF_TOKEN")
+        
+        # Configure RTC settings
+        rtc_config = {
+            "iceServers": [
+                {"urls": "stun:stun.l.google.com:19302"},
+                {"urls": "stun:stun1.l.google.com:19302"}
+            ]
+        }
+        
+        # Create FastRTC stream
+        stream = Stream(
+            handler=audio_handler,
+            rtc_configuration=rtc_config,
+            modality="audio",
+            mode="receive"  # We only receive audio, don't send
+        )
+        
+        # Mount the stream
+        app.mount("/stream", stream)
+        print("✅ FastRTC stream configured successfully!")
+        return stream
+        
+    except Exception as e:
+        print(f"⚠️ Warning: Failed to setup FastRTC stream: {e}")
+        print("Audio streaming may not work properly.")
+        return None
+
+
+# Main application setup
+def create_app(diarization_sys=None):
+    """Create the complete application"""
+    global diarization_system
+    
+    # Set the diarization system
+    if diarization_sys is not None:
+        diarization_system = diarization_sys
     
     # Create FastAPI app
     fastapi_app = create_fastapi_app()
@@ -1021,12 +1095,19 @@ def create_app():
     # Mount Gradio on FastAPI
     app = gr.mount_gradio_app(fastapi_app, gradio_interface, path="/")
     
+    # Setup FastRTC stream (will be called after initialization)
+    # Note: The stream setup happens when the system is initialized
+    
     return app, gradio_interface
 
 
 # Entry point for HuggingFace Spaces
 if __name__ == "__main__":
     try:
+        # Import your diarization system here
+        # from your_module import RealtimeSpeakerDiarization
+        # diarization_system = RealtimeSpeakerDiarization()
+        
         # Create the application
         app, interface = create_app()
         
@@ -1041,6 +1122,9 @@ if __name__ == "__main__":
         
     except Exception as e:
         print(f"Failed to launch application: {e}")
+        import traceback
+        traceback.print_exc()
+        
         # Fallback - launch just Gradio interface
         try:
             interface = create_interface()
@@ -1050,4 +1134,12 @@ if __name__ == "__main__":
                 share=False
             )
         except Exception as fallback_error:
-            print(f"Fallback launch also failed: {fallback_error}")
+            print(f"Fallback launch also failed: {fallback_error}")
+
+
+# Helper function to initialize with your diarization system
+def initialize_with_diarization_system(diarization_sys):
+    """Initialize the application with your diarization system"""
+    global diarization_system
+    diarization_system = diarization_sys
+    return create_app(diarization_sys)