Update app.py

app.py

@@ -30,6 +30,19 @@ import io
 from utils.utils import *
 from utils.markdown_utils import MarkdownConverter
 
+# Voice functionality imports
+import time
+import librosa
+from dataclasses import dataclass, field
+from pydub import AudioSegment
+try:
+    from voice_chat.utils.vad import get_speech_timestamps, collect_chunks, VadOptions
+    from voice_chat.gemma3n_inference import Gemma3nInference
+    VOICE_DEPENDENCIES_AVAILABLE = True
+except ImportError as e:
+    print(f"Voice dependencies not available: {e}")
+    VOICE_DEPENDENCIES_AVAILABLE = False
+
 # Math extension is optional for enhanced math rendering
 MATH_EXTENSION_AVAILABLE = False
 try:
@@ -38,6 +51,19 @@ try:
 except ImportError:
     pass
 
+# Initialize voice model early to avoid NameError
+voice_model = None
+if VOICE_DEPENDENCIES_AVAILABLE:
+    try:
+        print("Loading voice model...")
+        voice_model = Gemma3nInference(device='cuda' if torch.cuda.is_available() else 'cpu')
+        print("Warming up voice model...")
+        voice_model.warm_up()
+        print("✅ Voice model loaded and warmed up successfully")
+    except Exception as e:
+        print(f"⚠️ Voice model initialization failed: {e}")
+        voice_model = None
+
 
 class DOLPHIN:
     def __init__(self, model_id_or_path):
@@ -487,6 +513,149 @@ show_results_tab = False
 document_chunks = []
 document_embeddings = None
 
+# Voice chat parameters and state
+IN_CHANNELS = 1
+IN_RATE = 24000
+IN_CHUNK = 1024
+IN_SAMPLE_WIDTH = 2
+VAD_STRIDE = 0.5
+OUT_CHANNELS = 1
+OUT_RATE = 24000
+OUT_SAMPLE_WIDTH = 2
+OUT_CHUNK = 20 * 4096
+
+# Voice model already initialized earlier in the file
+
+@dataclass
+class VoiceAppState:
+    stream: np.ndarray | None = None
+    sampling_rate: int = 0
+    pause_detected: bool = False
+    started_talking: bool = False
+    stopped: bool = False
+    conversation: list = field(default_factory=list)
+
+
+# Voice functionality
+def run_vad(ori_audio, sr):
+    """Voice Activity Detection"""
+    _st = time.time()
+    try:
+        audio = ori_audio
+        if isinstance(audio, bytes):
+            audio = np.frombuffer(audio, dtype=np.int16)
+        audio = audio.astype(np.float32) / 32768.0
+        sampling_rate = 16000
+        if sr != sampling_rate:
+            audio = librosa.resample(audio, orig_sr=sr, target_sr=sampling_rate)
+
+        vad_parameters = {}
+        vad_parameters = VadOptions(**vad_parameters)
+        speech_chunks = get_speech_timestamps(audio, vad_parameters)
+        audio = collect_chunks(audio, speech_chunks)
+        duration_after_vad = audio.shape[0] / sampling_rate
+
+        if sr != sampling_rate:
+            vad_audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=sr)
+        else:
+            vad_audio = audio
+        vad_audio = np.round(vad_audio * 32768.0).astype(np.int16)
+        vad_audio_bytes = vad_audio.tobytes()
+
+        return duration_after_vad, vad_audio_bytes, round(time.time() - _st, 4)
+    except Exception as e:
+        msg = f"[asr vad error] audio_len: {len(ori_audio)/(sr*2):.3f} s, trace: {e}"
+        print(msg)
+        return -1, ori_audio, round(time.time() - _st, 4)
+
+def determine_pause(audio: np.ndarray, sampling_rate: int, state: VoiceAppState) -> bool:
+    """Determine if a pause happened in the audio stream"""
+    temp_audio = audio
+    dur_vad, _, time_vad = run_vad(temp_audio, sampling_rate)
+    duration = len(audio) / sampling_rate
+
+    if dur_vad > 0.5 and not state.started_talking:
+        print("started talking")
+        state.started_talking = True
+        return False
+
+    print(f"duration_after_vad: {dur_vad:.3f} s, time_vad: {time_vad:.3f} s")
+    return (duration - dur_vad) > 1
+
+def process_voice_audio(audio: tuple, state: VoiceAppState):
+    """Process streaming audio input"""
+    if not VOICE_DEPENDENCIES_AVAILABLE or voice_model is None:
+        return None, state
+        
+    if state.stream is None:
+        state.stream = audio[1]
+        state.sampling_rate = audio[0]
+    else:
+        state.stream = np.concatenate((state.stream, audio[1]))
+
+    pause_detected = determine_pause(state.stream, state.sampling_rate, state)
+    state.pause_detected = pause_detected
+
+    if state.pause_detected and state.started_talking:
+        return gr.Audio(recording=False), state
+    return None, state
+
+def generate_voice_response(state: VoiceAppState):
+    """Generate voice response from audio input"""
+    if not VOICE_DEPENDENCIES_AVAILABLE or voice_model is None:
+        return None, VoiceAppState()
+        
+    if not state.pause_detected and not state.started_talking:
+        return None, VoiceAppState()
+    
+    try:
+        audio_buffer = io.BytesIO()
+        segment = AudioSegment(
+            state.stream.tobytes(),
+            frame_rate=state.sampling_rate,
+            sample_width=state.stream.dtype.itemsize,
+            channels=(1 if len(state.stream.shape) == 1 else state.stream.shape[1]),
+        )
+        segment.export(audio_buffer, format="wav")
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            f.write(audio_buffer.getvalue())
+            temp_audio_path = f.name
+        
+        try:
+            # Generate text response from audio
+            text_response = voice_model.generate_response(temp_audio_path)
+            print(f"Generated voice response: {text_response}")
+            
+            # Convert text to speech
+            audio_response = voice_model.text_to_speech_simple(text_response)
+            
+            # Convert to format expected by Gradio
+            audio_segment = AudioSegment.from_file(io.BytesIO(audio_response), format="wav")
+            audio_array = np.array(audio_segment.get_array_of_samples())
+            
+            if audio_segment.channels == 2:
+                audio_array = audio_array.reshape((-1, 2))
+            
+            # Update conversation history
+            state.conversation.append({"role": "user", "content": {"path": temp_audio_path, "mime_type": "audio/wav"}})
+            state.conversation.append({"role": "assistant", "content": {"text": text_response}})
+            
+            return (audio_segment.frame_rate, audio_array), VoiceAppState(conversation=state.conversation)
+            
+        finally:
+            if os.path.exists(temp_audio_path):
+                os.unlink(temp_audio_path)
+                
+    except Exception as e:
+        print(f"Error generating voice response: {e}")
+        return None, VoiceAppState()
+
+def start_voice_recording(state: VoiceAppState):
+    """Start recording user voice input"""
+    if not state.stopped:
+        return gr.Audio(recording=True)
+    return gr.Audio(recording=False)
 
 def chunk_document(text, chunk_size=1024, overlap=100):
     """Split document into overlapping chunks for RAG"""
@@ -645,15 +814,22 @@ with gr.Blocks(
         # Home Tab
         with gr.TabItem("🏠 Home", id="home"):
             embedding_status = "✅ RAG ready" if embedding_model else "❌ RAG not loaded"
+            voice_status = "✅ Voice chat ready" if VOICE_DEPENDENCIES_AVAILABLE and voice_model else "❌ Voice chat not available"
             gr.Markdown(
-                "# Scholar Express - Local Gemma 3n Version\n"
-                "### Upload a research paper to get a web-friendly version with AI-generated alt text for accessibility. Includes an AI chatbot powered by local Gemma 3n.\n"
+                "# Scholar Express - Local Gemma 3n Version with Voice\n"
+                "### Upload a research paper to get a web-friendly version with AI-generated alt text for accessibility. Includes an AI chatbot and voice chat powered by local Gemma 3n.\n"
                 f"**System:** {model_status}\n"
                 f"**RAG System:** {embedding_status}\n"
+                f"**Voice Chat:** {voice_status}\n"
                 f"**DOLPHIN:** Local model for PDF processing\n"
-                f"**Gemma 3n:** Local model for alt text generation and chat\n"
+                f"**Gemma 3n:** Local model for alt text generation, chat, and voice\n"
                 f"**Alt Text:** Gemma 3n generates descriptive alt text for images\n"
-                f"**GPU:** {'CUDA available' if torch.cuda.is_available() else 'CPU only'}"
+                f"**GPU:** {'CUDA available' if torch.cuda.is_available() else 'CPU only'}\n\n"
+                "**Features:**\n"
+                "- 📄 PDF processing with OCR and layout analysis\n"
+                "- 💬 Text-based chat about your documents\n"
+                "- 🎙️ Voice chat with Gemma 3n (new!)\n"
+                "- ♿ AI-generated alt text for accessibility"
             )
             
             with gr.Column(elem_classes="upload-container"):
@@ -727,6 +903,51 @@ with gr.Blocks(
                 "*Ask questions about your processed document. The AI uses RAG (Retrieval-Augmented Generation) with local Gemma 3n to find relevant sections and provide accurate answers.*",
                 elem_id="chat-notice"
             )
+        
+        # Voice Chat Tab
+        with gr.TabItem("🎙️ Talk with Gemma", id="voice") as voice_tab:
+            voice_status = "✅ Voice chat ready" if VOICE_DEPENDENCIES_AVAILABLE and voice_model else "❌ Voice chat not available"
+            gr.Markdown(f"## Voice Chat with Gemma 3n\n{voice_status}")
+            
+            if VOICE_DEPENDENCIES_AVAILABLE and voice_model:
+                with gr.Row():
+                    with gr.Column():
+                        voice_input_audio = gr.Audio(
+                            label="Speak to Gemma", 
+                            sources=["microphone"], 
+                            type="numpy",
+                            streaming=True
+                        )
+                    with gr.Column():
+                        voice_output_audio = gr.Audio(
+                            label="Gemma's Response", 
+                            streaming=True, 
+                            autoplay=True
+                        )
+                        voice_chatbot = gr.Chatbot(
+                            label="Voice Conversation", 
+                            type="messages",
+                            height=300
+                        )
+                
+                with gr.Row():
+                    voice_stop_btn = gr.Button("⏹️ Stop Conversation", variant="stop")
+                
+                gr.Markdown(
+                    "*Speak naturally to Gemma 3n. The AI will listen to your voice, process your speech, and respond with both text and voice. You can have conversations before or after processing PDFs.*"
+                )
+                
+                # Voice state
+                voice_state = gr.State(value=VoiceAppState())
+            else:
+                gr.Markdown(
+                    "### Voice chat is not available\n"
+                    "To enable voice chat, please install the required dependencies:\n"
+                    "```bash\n"
+                    "pip install librosa pydub onnxruntime\n"
+                    "```\n"
+                    "And ensure the voice_chat directory is properly set up."
+                )
     
     # Event handlers
     process_btn.click(
@@ -814,6 +1035,46 @@ Please provide a clear and helpful answer based on the context provided."""
         lambda: "",
         outputs=[msg_input]
     )
+    
+    # Voice chat event handlers
+    if VOICE_DEPENDENCIES_AVAILABLE and voice_model:
+        # Stream processing
+        voice_stream = voice_input_audio.stream(
+            process_voice_audio,
+            [voice_input_audio, voice_state],
+            [voice_input_audio, voice_state],
+            stream_every=0.50,
+            time_limit=30,
+        )
+        
+        # Response generation
+        voice_respond = voice_input_audio.stop_recording(
+            generate_voice_response,
+            [voice_state],
+            [voice_output_audio, voice_state]
+        )
+        
+        # Update chatbot display
+        voice_respond.then(
+            lambda s: s.conversation, 
+            [voice_state], 
+            [voice_chatbot]
+        )
+        
+        # Restart recording
+        voice_restart = voice_output_audio.stop(
+            start_voice_recording,
+            [voice_state],
+            [voice_input_audio]
+        )
+        
+        # Stop conversation
+        voice_stop_btn.click(
+            lambda: (VoiceAppState(stopped=True), gr.Audio(recording=False)), 
+            None,
+            [voice_state, voice_input_audio], 
+            cancels=[voice_respond, voice_restart]
+        )
 
 
 if __name__ == "__main__":