Fix: send 'FINISH' text over WebSocket on stop to match server

app.py

@@ -1,415 +1,95 @@
-import asyncio
-import logging
-from fastapi import FastAPI, UploadFile, File, WebSocket
-from fastapi.responses import JSONResponse, StreamingResponse
+import time
+
 import gradio as gr
+import librosa
 import numpy as np
-import io
-
-import server_wrapper
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-app = FastAPI(title="SimulStreaming ASR")
-
-
-@app.on_event("startup")
-async def startup_event():
-    logger.info("Starting up... initializing model.")
-    loop = asyncio.get_event_loop()
-    try:
-        logger.info("Downloading Whisper model if not already present...")
-        await loop.run_in_executor(None, _ensure_model_downloaded)
-        logger.info("Model ready.")
-        
-        await loop.run_in_executor(None, server_wrapper.init_model)
-        logger.info("Model initialized successfully.")
-    except Exception as e:
-        logger.error(f"Error during model initialization: {e}")
-
-
-def _ensure_model_downloaded():
-    """Ensure the Whisper model is downloaded."""
-    import os
-    model_dir = os.path.expanduser('~/.cache/whisper')
-    model_path = os.path.join(model_dir, 'large-v3.pt')
-    
-    if not os.path.exists(model_path):
-        try:
-            logger.info(f"Downloading Whisper large-v3 model to {model_path}...")
-            import whisper
-            whisper.load_model('large-v3')
-            logger.info("Model downloaded successfully.")
-        except Exception as e:
-            logger.warning(f"Could not pre-download model: {e}")
+# import soundfile as sf
+from transformers import pipeline
+
+TARGET_SAMPLE_RATE = 16_000
+AUDIO_SECONDS_THRESHOLD = 2
+pipe = pipeline("audio-classification", model="MIT/ast-finetuned-audioset-10-10-0.4593")
+prediction = [{"score": 1, "label": "recording..."}]
+
+
+def normalize_waveform(waveform, datatype=np.float32):  # source datatype: np.int16
+    waveform = waveform.astype(dtype=datatype)
+    waveform /= 32768.0
+    return waveform
+
+
+def streaming_recording_fn(stream, new_chunk):
+    global prediction
+    sr, y = new_chunk
+    y = normalize_waveform(y)
+    y = librosa.resample(y, orig_sr=sr, target_sr=TARGET_SAMPLE_RATE)
+    if stream is not None:
+        if (stream.shape[-1] / TARGET_SAMPLE_RATE) >= AUDIO_SECONDS_THRESHOLD:
+            prediction = pipe(stream)
+            file_name = f'./audio/{time.strftime("%Y%m%d_%H%M%S", time.localtime())}.wav'
+            # # sf.write(file_name, stream, TARGET_SAMPLE_RATE)
+            print(f"SAVE AUDIO: {file_name}")
+            print(f">>>>>>1\t{y.shape=}, {stream.shape=}\n\t{prediction[0]=}")
+            stream = None
+        else:
+            stream = np.concatenate([stream, y], axis=-1)
+            print(f">>>>>>2\t{y.shape=}, {stream.shape=}")
     else:
-        logger.info(f"Model already present at {model_path}")
-
-
-@app.post("/api/reset")
-async def api_reset():
-    try:
-        server_wrapper.reset()
-        return JSONResponse({"status": "ok"})
-    except Exception as e:
-        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
-
-
-@app.post("/api/chunk")
-async def api_chunk(file: UploadFile = File(...)):
-    """Process a single audio chunk (streaming)."""
-    try:
-        raw = await file.read()
-        out = await asyncio.get_event_loop().run_in_executor(None, server_wrapper.process_chunk_from_bytes, raw)
-        return JSONResponse(out or {"text": ""})
-    except Exception as e:
-        logger.error(f"Error processing chunk: {e}")
-        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
-
+        stream = y
+        print(f">>>>>>3\t{y.shape=}, {stream.shape=}")
 
-@app.post("/api/finish")
-async def api_finish():
-    """Finish streaming and return final transcription."""
-    try:
-        out = await asyncio.get_event_loop().run_in_executor(None, server_wrapper.finish)
-        return JSONResponse(out or {"text": ""})
-    except Exception as e:
-        logger.error(f"Error finishing: {e}")
-        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
+    return stream, {i['label']: i['score'] for i in prediction}
 
 
-@app.websocket("/ws/audio")
-async def websocket_audio(websocket: WebSocket):
-    """WebSocket endpoint for real-time audio streaming."""
-    await websocket.accept()
-    logger.info("WebSocket connection established")
-    
-    try:
-        server_wrapper.reset()
-        
-        while True:
-            # Accept either binary frames (audio) or text frames (control messages like FINISH)
-            message = await websocket.receive()
-            data = None
-            is_text = False
-            if 'bytes' in message and message['bytes'] is not None:
-                data = message['bytes']
-            elif 'text' in message and message['text'] is not None:
-                data = message['text']
-                is_text = True
+def microphone_fn(waveform):
+    print('-' * 120)
+    print(f"{waveform=}")
+    sr, y = waveform
+    y = normalize_waveform(y)
+    y = librosa.resample(y, orig_sr=sr, target_sr=TARGET_SAMPLE_RATE)
+    result = pipe(y)
+    file_name = f'./audio/{time.strftime("%Y%m%d_%H%M%S", time.localtime())}.wav'
+    # sf.write(file_name, y, TARGET_SAMPLE_RATE)
+    return {i['label']: i['score'] for i in result}
 
-            if is_text:
-                # Control messages
-                if data == "FINISH":
-                    result = await asyncio.get_event_loop().run_in_executor(None, server_wrapper.finish)
-                    await websocket.send_json({"type": "finish", **(result or {})})
-                    break
-                elif data == "RESET":
-                    server_wrapper.reset()
-                    await websocket.send_json({"type": "reset", "status": "ok"})
-                else:
-                    # Unknown text message - ignore or log
-                    logger.debug(f"Unknown WS text message: {data}")
-            else:
-                # Binary audio chunk (or binary control marker)
-                try:
-                    # If client sent the 4-byte control marker 0xFF 0xFF 0xFF 0xFF, treat as FINISH
-                    if isinstance(data, (bytes, bytearray)) and data == b"\xFF\xFF\xFF\xFF":
-                        result = await asyncio.get_event_loop().run_in_executor(None, server_wrapper.finish)
-                        await websocket.send_json({"type": "finish", **(result or {})})
-                        break
 
-                    result = await asyncio.get_event_loop().run_in_executor(None, server_wrapper.process_chunk_from_bytes, data)
-                    if result and result.get("text"):
-                        await websocket.send_json({"type": "update", "text": result["text"]})
-                except Exception as e:
-                    logger.error(f"Error processing audio chunk via websocket: {e}")
-                    await websocket.send_json({"type": "error", "message": str(e)})
-    
-    except Exception as e:
-        logger.error(f"WebSocket error: {e}")
-        try:
-            await websocket.send_json({"type": "error", "message": str(e)})
-        except:
-            pass
-    finally:
-        await websocket.close()
-        logger.info("WebSocket connection closed")
+def file_fn(waveform):
+    print('-' * 120)
+    print(f"{waveform=}")
+    sr, y = waveform
+    y = normalize_waveform(y)
+    y = librosa.resample(y, orig_sr=sr, target_sr=TARGET_SAMPLE_RATE)
+    result = pipe(y)
+    file_name = f'./audio/{time.strftime("%Y%m%d_%H%M%S", time.localtime())}.wav'
+    # sf.write(file_name, y, TARGET_SAMPLE_RATE)
+    return {i['label']: i['score'] for i in result}
 
 
-def create_ui():
-    with gr.Blocks(title="Streaming ASR", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🎙️ Streaming ASR — SimulWhisper
-        
-        Transcripción en tiempo real mientras hablas.
-        
-        **Instrucciones:**
-        1. Haz clic en **"🔴 Start Recording"** 
-        2. Habla naturalmente - verás la transcripción EN TIEMPO REAL
-        3. Haz clic en **"⏹️ Stop Recording"** cuando termines
-        """)
-        
-        with gr.Row():
-            start_btn = gr.Button("🔴 Start Recording", size="lg", variant="primary", scale=1)
-            stop_btn = gr.Button("⏹️ Stop Recording", size="lg", variant="stop", scale=1)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("### Status")
-                status = gr.Textbox(
-                    value="Ready",
-                    interactive=False,
-                    show_label=False,
-                    lines=2
-                )
-            
-            with gr.Column(scale=2):
-                gr.Markdown("### 📝 Transcripción en Vivo")
-                transcript = gr.Textbox(
-                    show_label=False,
-                    lines=8,
-                    interactive=False,
-                    placeholder="La transcripción aparecerá aquí en tiempo real..."
-                )
-        
-        # JavaScript para captura real-time con WebSocket
-        html_js = """
-<script>
-let mediaRecorder;
-let audioCtx;
-let source;
-let processor;
-let recording = false;
-let ws = null;
-let chunkSize = 16000 * 0.5; // 0.5 seconds at 16kHz
-let startBtn = null;
-let stopBtn = null;
-let statusDiv = null;
-let transcriptDiv = null;
+streaming_demo = gr.Interface(
+    fn=streaming_recording_fn,
+    inputs=["state", gr.Audio(sources=["microphone"], streaming=True)],
+    outputs=["state", "label"],
+    live=True,
+)
 
-function to16BitPCM(float32Array) {
-    const l = float32Array.length;
-    const buffer = new ArrayBuffer(l * 2);
-    const view = new DataView(buffer);
-    let offset = 0;
-    for (let i = 0; i < l; i++) {
-        let s = Math.max(-1, Math.min(1, float32Array[i]));
-        view.setInt16(offset, s < 0 ? s * 0x8000 : s * 0x7FFF, true);
-        offset += 2;
-    }
-    return buffer;
-}
 
-function writeWAV(samples, sampleRate) {
-    const buffer = new ArrayBuffer(44 + samples.byteLength);
-    const view = new DataView(buffer);
-    function writeString(view, offset, string) {
-        for (let i = 0; i < string.length; i++) {
-            view.setUint8(offset + i, string.charCodeAt(i));
-        }
-    }
-    writeString(view, 0, 'RIFF');
-    view.setUint32(4, 36 + samples.byteLength, true);
-    writeString(view, 8, 'WAVE');
-    writeString(view, 12, 'fmt ');
-    view.setUint32(16, 16, true);
-    view.setUint16(20, 1, true);
-    view.setUint16(22, 1, true);
-    view.setUint32(24, sampleRate, true);
-    view.setUint32(28, sampleRate * 2, true);
-    view.setUint16(32, 2, true);
-    view.setUint16(34, 16, true);
-    writeString(view, 36, 'data');
-    view.setUint32(40, samples.byteLength, true);
-    const bytes = new Uint8Array(buffer, 44);
-    bytes.set(new Uint8Array(samples));
-    return buffer;
-}
 
-async function resampleAudio(float32Array, fromSampleRate, toSampleRate) {
-    if (fromSampleRate === toSampleRate) {
-        return float32Array;
-    }
-    const length = Math.round(float32Array.length * toSampleRate / fromSampleRate);
-    const offlineCtx = new OfflineAudioContext(1, length, toSampleRate);
-    const buffer = offlineCtx.createBuffer(1, float32Array.length, fromSampleRate);
-    buffer.copyToChannel(float32Array, 0, 0);
-    const src = offlineCtx.createBufferSource();
-    src.buffer = buffer;
-    src.connect(offlineCtx.destination);
-    src.start(0);
-    const rendered = await offlineCtx.startRendering();
-    return rendered.getChannelData(0);
-}
+with gr.Blocks() as example:
+    inputs = [gr.Audio(sources=["upload"], type="numpy")]
+    output = gr.Label()
 
-async function sendChunk(float32Array, sampleRate) {
-    if (!ws || ws.readyState !== WebSocket.OPEN) return;
-    
-    try {
-        let resampled = await resampleAudio(float32Array, sampleRate, 16000);
-        const pcm16 = to16BitPCM(resampled);
-        const wav = writeWAV(pcm16, 16000);
-        ws.send(wav);
-    } catch (e) {
-        console.error('Error sending chunk:', e);
-    }
-}
-
-async function startRecording() {
-    try {
-        if (recording) return;
-        
-        console.log('Starting recording...');
-        
-        // Connect WebSocket
-        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-        ws = new WebSocket(protocol + '//' + window.location.host + '/ws/audio');
-        
-        ws.onopen = () => {
-            console.log('WebSocket connected');
-            updateStatus('🔴 Recording... listening');
-        };
-        
-        ws.onmessage = (event) => {
-            const data = JSON.parse(event.data);
-            if (data.type === 'update' && data.text) {
-                updateTranscript(data.text);
-            } else if (data.type === 'finish') {
-                console.log('Transcription finished:', data);
-                updateStatus('✅ Done');
-            }
-        };
-        
-        ws.onerror = (error) => {
-            console.error('WebSocket error:', error);
-            updateStatus('❌ Connection error');
-        };
-        
-        ws.onclose = () => {
-            console.log('WebSocket closed');
-            recording = false;
-        };
-        
-        // Start audio capture
-        recording = true;
-        audioCtx = new (window.AudioContext || window.webkitAudioContext)();
-        const stream = await navigator.mediaDevices.getUserMedia({ 
-            audio: { 
-                echoCancellation: false, 
-                noiseSuppression: false,
-                autoGainControl: false
-            } 
-        });
-        source = audioCtx.createMediaStreamSource(stream);
-        processor = audioCtx.createScriptProcessor(4096, 1, 1);
-        
-        let buffer = [];
-        
-        processor.onaudioprocess = function(e) {
-            const ch = e.inputBuffer.getChannelData(0);
-            for (let i = 0; i < ch.length; i++) {
-                buffer.push(ch[i]);
-            }
-            
-            // Send chunk every 0.5 seconds
-            if (buffer.length >= chunkSize) {
-                const chunk = new Float32Array(buffer.slice(0, chunkSize));
-                buffer = buffer.slice(chunkSize);
-                sendChunk(chunk, audioCtx.sampleRate);
-            }
-        };
-        
-        source.connect(processor);
-        processor.connect(audioCtx.destination);
-        
-    } catch (e) {
-        console.error('Error starting recording:', e);
-        updateStatus('❌ Error: ' + e.message);
-        recording = false;
-    }
-}
-
-function stopRecording() {
-    if (!recording) return;
-    
-    recording = false;
-    updateStatus('⏹️ Stopping...');
-    
-    if (source && source.mediaStream) {
-        const tracks = source.mediaStream.getTracks();
-        tracks.forEach(t => t.stop());
-    }
-    if (processor) processor.disconnect();
-    if (source) source.disconnect();
-    
-    // Send finish signal (binary marker) so server recognizes it
-    if (ws && ws.readyState === WebSocket.OPEN) {
-        ws.send(new Uint8Array([0xFF, 0xFF, 0xFF, 0xFF]));
-        setTimeout(() => {
-            if (ws) ws.close();
-        }, 500);
-    }
-}
-
-function updateTranscript(text) {
-    const textareas = document.querySelectorAll('textarea');
-    if (textareas.length >= 2) {
-        textareas[1].value = text;
-        textareas[1].dispatchEvent(new Event('input', { bubbles: true }));
-    }
-}
-
-function updateStatus(text) {
-    const textareas = document.querySelectorAll('textarea');
-    if (textareas.length >= 1) {
-        textareas[0].value = text;
-        textareas[0].dispatchEvent(new Event('input', { bubbles: true }));
-    }
-}
-
-// Find and attach button listeners
-function attachButtons() {
-    const buttons = document.querySelectorAll('button');
-    console.log('Found ' + buttons.length + ' buttons');
-    
-    if (buttons.length >= 2) {
-        startBtn = buttons[0];
-        stopBtn = buttons[1];
-        
-        startBtn.addEventListener('click', startRecording);
-        stopBtn.addEventListener('click', stopRecording);
-        
-        console.log('Buttons attached successfully');
-    }
-}
-
-// Try to attach buttons when page loads
-document.addEventListener('DOMContentLoaded', () => {
-    console.log('DOM loaded');
-    setTimeout(attachButtons, 1000);
-});
-
-// Also try immediately
-setTimeout(attachButtons, 500);
-</script>
-"""
-        
-        gr.HTML(html_js)
-    
-    return demo
-
-
-demo = create_ui()
-
-# Mount Gradio app on FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
+    examples = [
+        ["audio/cantina.wav"],
+        ["audio/cat.mp3"]
+    ]
+    ex = gr.Examples(examples,
+                     fn=file_fn, inputs=inputs, outputs=output,
+                     run_on_click=True)
 
+with gr.Blocks() as demo:
+    gr.TabbedInterface([streaming_demo],
+                       ["Streaming"])
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-
-
 
+    demo.launch(share=True)