fix: enhance audio processing in transcribe function with improved buffering and chunk handling

app.py

@@ -24,50 +24,105 @@ def load_model():
     return model
 
 @spaces.GPU(duration=120)
-def transcribe(audio, state=""):
+def transcribe(audio, state="", audio_buffer=None, last_processed_time=0):
     # Load the model inside the GPU worker process
     import numpy as np
     import soundfile as sf
     import librosa
     import os
-    model = load_model()
+    if audio_buffer is None:
+        audio_buffer = []
     
     if audio is None or isinstance(audio, int):
         print(f"Skipping invalid audio input: {type(audio)}")
-        return state, state
+        return state, state, audio_buffer, last_processed_time
+    
     print(f"Received audio input of type: {type(audio)}")
-    print(f"Audio shape: {audio.shape if isinstance(audio, np.ndarray) else 'N/A'}")
     
     if isinstance(audio, tuple) and len(audio) == 2 and isinstance(audio[1], np.ndarray):
-        # Handle tuple of (sample_rate, audio_array)
-        print(f"Tuple contents: {audio}")
         sample_rate, audio_data = audio
+        print(f"Sample rate: {sample_rate}, Audio shape: {audio_data.shape}")
+        
+        # Append chunk to buffer
+        audio_buffer.append(audio_data)
+        
+        # Calculate total duration in seconds
+        total_samples = sum(arr.shape[0] for arr in audio_buffer)
+        total_duration = total_samples / sample_rate
+        print(f"Total buffered duration: {total_duration:.2f}s")
+        
+        # Process 3-second chunks with 1-second step size (2-second overlap)
+        chunk_duration = 3.0  # seconds
+        step_size = 1.0      # seconds
+        min_samples = int(chunk_duration * 16000)  # 3s at 16kHz
+        
+        if total_duration < chunk_duration:
+            print(f"Buffering audio, total duration: {total_duration:.2f}s")
+            return state, state, audio_buffer, last_processed_time
+        
         try:
-            # Resample to 16kHz for NeMo
+            # Concatenate buffered chunks
+            full_audio = np.concatenate(audio_buffer)
+            
+            # Resample to 16kHz if needed
             if sample_rate != 16000:
                 print(f"Resampling from {sample_rate}Hz to 16000Hz")
-                audio_data = librosa.resample(audio_data.astype(float), orig_sr=sample_rate, target_sr=16000)
-            # Save to temporary WAV file
-            temp_file = "temp_audio.wav"
-            sf.write(temp_file, audio_data, samplerate=16000)
-            print(f"Processing temporary audio file: {temp_file}")
+                full_audio = librosa.resample(full_audio.astype(float), orig_sr=sample_rate, target_sr=16000)
+                sample_rate = 16000
+            else:
+                full_audio = full_audio.astype(float)
+            
+            # Process 3-second chunks
+            new_state = state
+            current_time = last_processed_time
+            total_samples_16k = len(full_audio)
+            
+            while current_time + chunk_duration <= total_duration:
+                start_sample = int(current_time * sample_rate)
+                end_sample = int((current_time + chunk_duration) * sample_rate)
+                if end_sample > total_samples_16k:
+                    break
+                
+                chunk = full_audio[start_sample:end_sample]
+                print(f"Processing chunk from {current_time:.2f}s to {current_time + chunk_duration:.2f}s")
+                
+                # Save to temporary WAV file
+                temp_file = "temp_audio.wav"
+                sf.write(temp_file, chunk, samplerate=16000)
+                
+                # Transcribe
+                hypothesis = model.transcribe([temp_file])[0]
+                transcription = hypothesis.text
+                print(f"Transcription: {transcription}")
+                
+                os.remove(temp_file)
+                print("Temporary file removed.")
+                
+                # Append transcription if non-empty
+                if transcription.strip():
+                    new_state = new_state + " " + transcription if new_state else transcription
+                
+                current_time += step_size
+            
+            # Update last processed time
+            last_processed_time = current_time
+            
+            # Trim buffer to keep only unprocessed audio
+            keep_samples = int((total_duration - current_time) * sample_rate)
+            if keep_samples > 0:
+                audio_buffer = [full_audio[-keep_samples:]]
+            else:
+                audio_buffer = []
             
-            # Transcribe and extract only the text (string)
-            hypothesis = model.transcribe([temp_file])[0]
-            print(f"Hypothesis: {hypothesis}")
-            transcription = hypothesis.text  # Extract the text attribute (string)
-            print(f"Transcription: {transcription}")
-                    
-            os.remove(temp_file)  # Clean up
-            print("Temporary file removed.")
+            print(f"New state: {new_state}")
+            return new_state, new_state, audio_buffer, last_processed_time
+        
         except Exception as e:
             print(f"Error processing audio: {e}")
-            return state, state
-
-        new_state = state + " " + transcription if state else transcription
-        print(f"New state: {new_state}")
-        return new_state, new_state
-    return state, state
+            return state, state, audio_buffer, last_processed_time
+    
+    print(f"Invalid audio input format: {type(audio)}")
+    return state, state, audio_buffer, last_processed_time
 
 # Define the Gradio interface
 with gr.Blocks(title="Real-time Speech-to-Text with NeMo") as demo:
@@ -98,22 +153,23 @@ with gr.Blocks(title="Real-time Speech-to-Text with NeMo") as demo:
     
     # State to store the ongoing transcription
     state = gr.State("")
-    
+    audio_buffer = gr.State(value=None)
+    last_processed_time = gr.State(value=0)
     # Handle the audio stream
     audio_input.stream(
         fn=transcribe,
-        inputs=[audio_input, state],
-        outputs=[state, streaming_text],
+        inputs=[audio_input, state, audio_buffer, last_processed_time],
+        outputs=[state, streaming_text, audio_buffer, last_processed_time],
     )
 
     # Clear the transcription
     def clear_transcription():
-        return "", "", ""
-    
+        return "", "", None, 0
+
     clear_btn.click(
         fn=clear_transcription,
         inputs=[],
-        outputs=[text_output, streaming_text, state]
+        outputs=[text_output, streaming_text, audio_buffer, last_processed_time]
     )
     
     # Update the main text output when the state changes