complete change

client.py

@@ -2,80 +2,173 @@ import asyncio
 import websockets
 import requests
 import ssl
+import wave
+import logging
+import sys
 
 # Parameters for reading and sending the audio
 AUDIO_FILE_URL = "https://raw.githubusercontent.com/AshDavid12/runpod-serverless-forked/main/test_hebrew.wav"  # Use WAV file
 #AUDIO_FILE_URL = "https://raw.githubusercontent.com/AshDavid12/hugging_face_ivrit_streaming/main/long_hebrew.wav"
 
 
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s: %(message)s',
+                    handlers=[logging.StreamHandler(sys.stdout)], force=True)
+logger = logging.getLogger(__name__)
+
+async def send_receive():
+    uri = "wss://gigaverse-ivrit-ai-streaming.hf.space/ws"  # Update with your server's address if needed
+    logger.info(f"Connecting to server at {uri}")
+    try:
+        async with websockets.connect(uri) as websocket:
+            logger.info("WebSocket connection established")
+            # Start tasks for sending and receiving
+            send_task = asyncio.create_task(send_audio(websocket))
+            receive_task = asyncio.create_task(receive_transcriptions(websocket))
+            await asyncio.gather(send_task, receive_task)
+    except Exception as e:
+        logger.error(f"WebSocket connection error: {e}")
+
 async def send_audio(websocket):
-    buffer_size = 512 * 1024  #HAVE TO HAVE 512!!
-    audio_buffer = bytearray()
-
-    with requests.get(AUDIO_FILE_URL, stream=True, allow_redirects=False) as response:
-        if response.status_code == 200:
-            print("Starting to stream audio file...")
-
-            for chunk in response.iter_content(chunk_size=1024):  # Stream in chunks
-                if chunk:
-                    audio_buffer.extend(chunk)
-                    #print(f"Received audio chunk of size {len(chunk)} bytes.")
-
-                    # Send buffered audio data once it's large enough
-                    #if len(audio_buffer) >= buffer_size:
-                    await websocket.send(audio_buffer)
-                        #print(f"Sent {len(audio_buffer)} bytes of audio data.")
-                    audio_buffer.clear()
-                    await asyncio.sleep(0.01)
-
-            print("Finished sending audio.")
-        else:
-            print(f"Failed to download audio file. Status code: {response.status_code}")
-
-
-async def receive_transcription(websocket):
-    while True:
-        try:
-
-            transcription = await websocket.recv()
-            # Receive transcription from the server
-            print(f"Transcription: {transcription}")
-        except Exception as e:
-            print(f"Error receiving transcription: {e}")
-            #await asyncio.sleep(30)
-            break
-
-
-async def send_heartbeat(websocket):
-    while True:
-        try:
-            await websocket.ping()
-            print("Sent keepalive ping")
-        except websockets.ConnectionClosed:
-            print("Connection closed, stopping heartbeat")
-            break
-        await asyncio.sleep(30)  # Send ping every 30 seconds (adjust as needed)
-
-
-async def run_client():
-    uri = ("wss://gigaverse-ivrit-ai-streaming.hf.space/wtranscribe")  # WebSocket URL
-    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-    ssl_context.check_hostname = False
-    ssl_context.verify_mode = ssl.CERT_NONE
-    while True:
-        try:
-            async with websockets.connect(uri, ssl=ssl_context, ping_timeout=1000, ping_interval=50) as websocket:
-                await asyncio.gather(
-                    send_audio(websocket),
-                    receive_transcription(websocket),
-                    send_heartbeat(websocket)
-                )
-        except websockets.ConnectionClosedError as e:
-            print(f"WebSocket closed with error: {e}")
-        # except Exception as e:
-        #     print(f"Unexpected error: {e}")
-        #
-        # print("Reconnecting in 5 seconds...")
-        # await asyncio.sleep(5)  # Wait 5 seconds before reconnecting
-
-asyncio.run(run_client())
+    wav_file = 'path/to/your/audio.wav'  # Replace with the path to your WAV file
+    logger.info(f"Opening WAV file: {wav_file}")
+
+    try:
+        # Open the WAV file
+        wf = wave.open(wav_file, 'rb')
+
+        # Log WAV file parameters
+        channels = wf.getnchannels()
+        sampwidth = wf.getsampwidth()
+        framerate = wf.getframerate()
+        nframes = wf.getnframes()
+        duration = nframes / framerate
+        logger.debug(f"WAV file parameters: channels={channels}, sample_width={sampwidth}, framerate={framerate}, frames={nframes}, duration={duration:.2f}s")
+
+        # Ensure the WAV file has the expected parameters
+        if channels != 1 or sampwidth != 2 or framerate != 16000:
+            logger.error("WAV file must be mono channel, 16-bit samples, 16kHz sampling rate")
+            return
+
+        chunk_duration = 0.1  # in seconds
+        chunk_size = int(framerate * chunk_duration)
+        logger.info(f"Starting to send audio data in chunks of {chunk_duration}s ({chunk_size} frames)")
+
+        total_chunks = 0
+        total_bytes_sent = 0
+
+        while True:
+            data = wf.readframes(chunk_size)
+            if not data:
+                logger.info("End of WAV file reached")
+                break
+            await websocket.send(data)
+            total_chunks += 1
+            total_bytes_sent += len(data)
+            logger.debug(f"Sent chunk {total_chunks}: {len(data)} bytes")
+            await asyncio.sleep(chunk_duration)  # Simulate real-time streaming
+
+        logger.info(f"Finished sending audio data: {total_chunks} chunks sent, total bytes sent: {total_bytes_sent}")
+    except Exception as e:
+        logger.error(f"Send audio error: {e}")
+    finally:
+        wf.close()
+        logger.info("WAV file closed")
+
+async def receive_transcriptions(websocket):
+    try:
+        logger.info("Starting to receive transcriptions")
+        async for message in websocket:  # This is the same as websocket.recv()
+            logger.info(f"Received transcription: {message}")
+            print(f"Transcription: {message}")
+    except Exception as e:
+        logger.error(f"Receive transcription error: {e}")
+
+if __name__ == "__main__":
+    asyncio.run(send_receive())
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# async def send_audio(websocket):
+#     buffer_size = 512 * 1024  #HAVE TO HAVE 512!!
+#     audio_buffer = bytearray()
+#
+#     with requests.get(AUDIO_FILE_URL, stream=True, allow_redirects=False) as response:
+#         if response.status_code == 200:
+#             print("Starting to stream audio file...")
+#
+#             for chunk in response.iter_content(chunk_size=1024):  # Stream in chunks
+#                 if chunk:
+#                     audio_buffer.extend(chunk)
+#                     #print(f"Received audio chunk of size {len(chunk)} bytes.")
+#
+#                     # Send buffered audio data once it's large enough
+#                 if len(audio_buffer) >= buffer_size:
+#                     await websocket.send(audio_buffer)
+#                         #print(f"Sent {len(audio_buffer)} bytes of audio data.")
+#                     audio_buffer.clear()
+#                     await asyncio.sleep(0.01)
+#
+#             print("Finished sending audio.")
+#         else:
+#             print(f"Failed to download audio file. Status code: {response.status_code}")
+#
+#
+# async def receive_transcription(websocket):
+#     while True:
+#         try:
+#
+#             transcription = await websocket.recv()
+#             # Receive transcription from the server
+#             print(f"Transcription: {transcription}")
+#         except Exception as e:
+#             print(f"Error receiving transcription: {e}")
+#             #await asyncio.sleep(30)
+#             break
+#
+#
+# async def send_heartbeat(websocket):
+#     while True:
+#         try:
+#             await websocket.ping()
+#             print("Sent keepalive ping")
+#         except websockets.ConnectionClosed:
+#             print("Connection closed, stopping heartbeat")
+#             break
+#         await asyncio.sleep(30)  # Send ping every 30 seconds (adjust as needed)
+#
+#
+# async def run_client():
+#     uri = ("wss://gigaverse-ivrit-ai-streaming.hf.space/wtranscribe")  # WebSocket URL
+#     ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+#     ssl_context.check_hostname = False
+#     ssl_context.verify_mode = ssl.CERT_NONE
+#     while True:
+#         try:
+#             async with websockets.connect(uri, ssl=ssl_context, ping_timeout=1000, ping_interval=50) as websocket:
+#                 await asyncio.gather(
+#                     send_audio(websocket),
+#                     receive_transcription(websocket),
+#                     send_heartbeat(websocket)
+#                 )
+#         except websockets.ConnectionClosedError as e:
+#             print(f"WebSocket closed with error: {e}")
+#         # except Exception as e:
+#         #     print(f"Unexpected error: {e}")
+#         #
+#         # print("Reconnecting in 5 seconds...")
+#         # await asyncio.sleep(5)  # Wait 5 seconds before reconnecting
+#
+# asyncio.run(run_client())

infer.py

@@ -1,6 +1,8 @@
 import base64
 import faster_whisper
 import tempfile
+
+import numpy as np
 import torch
 import time
 import requests
@@ -15,6 +17,7 @@ import asyncio
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s',
                     handlers=[logging.StreamHandler(sys.stdout)], force=True)
+logger = logging.getLogger(__name__)
 #logging.getLogger("asyncio").setLevel(logging.DEBUG)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 logging.info(f'Device selected: {device}')
@@ -31,49 +34,49 @@ logging.info(f'Max payload size set to: {MAX_PAYLOAD_SIZE} bytes')
 app = FastAPI()
 
 
-class InputData(BaseModel):
-    type: str
-    data: Optional[str] = None  # Used for blob input
-    url: Optional[str] = None  # Used for url input
-
-
-def download_file(url, max_size_bytes, output_filename, api_key=None):
-    """
-    Download a file from a given URL with size limit and optional API key.
-    """
-    logging.debug(f'Starting file download from URL: {url}')
-    try:
-        headers = {}
-        if api_key:
-            headers['Authorization'] = f'Bearer {api_key}'
-            logging.debug('API key provided, added to headers')
-
-        response = requests.get(url, stream=True, headers=headers)
-        response.raise_for_status()
-
-        file_size = int(response.headers.get('Content-Length', 0))
-        logging.info(f'File size: {file_size} bytes')
-
-        if file_size > max_size_bytes:
-            logging.error(f'File size exceeds limit: {file_size} > {max_size_bytes}')
-            return False
-
-        downloaded_size = 0
-        with open(output_filename, 'wb') as file:
-            for chunk in response.iter_content(chunk_size=8192):
-                downloaded_size += len(chunk)
-                logging.debug(f'Downloaded {downloaded_size} bytes')
-                if downloaded_size > max_size_bytes:
-                    logging.error('Downloaded size exceeds maximum allowed payload size')
-                    return False
-                file.write(chunk)
-
-        logging.info(f'File downloaded successfully: {output_filename}')
-        return True
-
-    except requests.RequestException as e:
-        logging.error(f"Error downloading file: {e}")
-        return False
+# class InputData(BaseModel):
+#     type: str
+#     data: Optional[str] = None  # Used for blob input
+#     url: Optional[str] = None  # Used for url input
+#
+#
+# def download_file(url, max_size_bytes, output_filename, api_key=None):
+#     """
+#     Download a file from a given URL with size limit and optional API key.
+#     """
+#     logging.debug(f'Starting file download from URL: {url}')
+#     try:
+#         headers = {}
+#         if api_key:
+#             headers['Authorization'] = f'Bearer {api_key}'
+#             logging.debug('API key provided, added to headers')
+#
+#         response = requests.get(url, stream=True, headers=headers)
+#         response.raise_for_status()
+#
+#         file_size = int(response.headers.get('Content-Length', 0))
+#         logging.info(f'File size: {file_size} bytes')
+#
+#         if file_size > max_size_bytes:
+#             logging.error(f'File size exceeds limit: {file_size} > {max_size_bytes}')
+#             return False
+#
+#         downloaded_size = 0
+#         with open(output_filename, 'wb') as file:
+#             for chunk in response.iter_content(chunk_size=8192):
+#                 downloaded_size += len(chunk)
+#                 logging.debug(f'Downloaded {downloaded_size} bytes')
+#                 if downloaded_size > max_size_bytes:
+#                     logging.error('Downloaded size exceeds maximum allowed payload size')
+#                     return False
+#                 file.write(chunk)
+#
+#         logging.info(f'File downloaded successfully: {output_filename}')
+#         return True
+#
+#     except requests.RequestException as e:
+#         logging.error(f"Error downloading file: {e}")
+#         return False
 
 
 @app.get("/")
@@ -81,97 +84,208 @@ async def read_root():
     return {"message": "This is the Ivrit AI Streaming service."}
 
 
-async def transcribe_core_ws(audio_file):
-    """
-    Transcribe the audio file and return only the segments that have not been processed yet.
+# async def transcribe_core_ws(audio_file):
+#     ret = {'segments': []}
+#
+#     try:
+#
+#         logging.debug(f"Initiating model transcription for file: {audio_file}")
+#
+#         segs, _ = await asyncio.to_thread(model.transcribe, audio_file, language='he', word_timestamps=True)
+#         logging.info('Transcription completed successfully.')
+#     except Exception as e:
+#         logging.error(f"Error during transcription: {e}")
+#         raise e
+#
+#     # Track the new segments and update the last transcribed time
+#     for s in segs:
+#         logging.info(f"Processing segment with start time: {s.start} and end time: {s.end}")
+#
+#         # Only process segments that start after the last transcribed time
+#         logging.info(f"New segment found starting at {s.start} seconds.")
+#         words = [{'start': w.start, 'end': w.end, 'word': w.word, 'probability': w.probability} for w in s.words]
+#
+#         seg = {
+#             'id': s.id, 'seek': s.seek, 'start': s.start, 'end': s.end, 'text': s.text,
+#             'avg_logprob': s.avg_logprob, 'compression_ratio': s.compression_ratio,
+#             'no_speech_prob': s.no_speech_prob, 'words': words
+#         }
+#         logging.info(f'Adding new transcription segment: {seg}')
+#         ret['segements'].append(seg)
+#
+#         # Update the last transcribed time to the end of the current segment
+#
+#
+# #logging.info(f"Returning {len(ret['new_segments'])} new segments and updated last transcribed time.")
+#     return ret
+
+
+import tempfile
 
-    :param audio_file: Path to the growing audio file.
-    :param last_transcribed_time: The last time (in seconds) that was transcribed.
-    :return: Newly transcribed segments and the updated last transcribed time.
-    """
-    ret = {'segments': []}
 
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """WebSocket endpoint to handle client connections."""
+    await websocket.accept()
+    client_ip = websocket.client.host
+    logger.info(f"Client connected: {client_ip}")
+    try:
+        await process_audio_stream(websocket)
+    except WebSocketDisconnect:
+        logger.info(f"Client disconnected: {client_ip}")
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        await websocket.close()
+
+async def process_audio_stream(websocket: WebSocket):
+    """Continuously receive audio chunks and initiate transcription tasks."""
+    sampling_rate = 16000
+    min_chunk_size = 1  # in seconds
+    audio_buffer = np.array([], dtype=np.float32)
+
+    transcription_task = None
+    chunk_counter = 0
+    total_bytes_received = 0
+
+    while True:
+        try:
+            # Receive audio data from client
+            data = await websocket.receive_bytes()
+            if not data:
+                logger.info("No data received, closing connection")
+                break
+            chunk_counter += 1
+            chunk_size = len(data)
+            total_bytes_received += chunk_size
+            logger.debug(f"Received chunk {chunk_counter}: {chunk_size} bytes")
+
+            audio_chunk = process_received_audio(data)
+            logger.debug(f"Processed audio chunk {chunk_counter}: {len(audio_chunk)} samples")
+
+            audio_buffer = np.concatenate((audio_buffer, audio_chunk))
+            logger.debug(f"Audio buffer size: {len(audio_buffer)} samples")
+        except Exception as e:
+            logger.error(f"Error receiving data: {e}")
+            break
+
+        # Check if enough audio has been buffered
+        if len(audio_buffer) >= min_chunk_size * sampling_rate:
+            if transcription_task is None or transcription_task.done():
+                # Start a new transcription task
+                logger.info(f"Starting transcription task for {len(audio_buffer)} samples")
+                transcription_task = asyncio.create_task(
+                    transcribe_and_send(websocket, audio_buffer.copy())
+                )
+                audio_buffer = np.array([], dtype=np.float32)
+                logger.debug("Audio buffer reset after starting transcription task")
+
+async def transcribe_and_send(websocket: WebSocket, audio_data):
+    """Run transcription in a separate thread and send the result to the client."""
+    logger.debug(f"Transcription task started for {len(audio_data)} samples")
+    transcription_result = await asyncio.to_thread(sync_transcribe_audio, audio_data)
+    if transcription_result:
+        try:
+            # Send the result as JSON
+            await websocket.send_json(transcription_result)
+            logger.info("Transcription JSON sent to client")
+        except Exception as e:
+            logger.error(f"Error sending transcription: {e}")
+    else:
+        logger.warning("No transcription result to send")
+
+def sync_transcribe_audio(audio_data):
+    """Synchronously transcribe audio data using the ASR model and format the result."""
     try:
-        # Transcribe the entire audio file
-        logging.debug(f"Initiating model transcription for file: {audio_file}")
 
-        segs, _ = await asyncio.to_thread(model.transcribe, audio_file, language='he', word_timestamps=True)
-        logging.info('Transcription completed successfully.')
+        logger.info('Starting transcription...')
+        segments, info = model.transcribe(
+            audio_data, language="he", beam_size=5, word_timestamps=True
+        )
+        logger.info('Transcription completed')
+
+        # Build the transcription result as per your requirement
+        ret = {'segments': []}
+
+        for s in segments:
+            logger.debug(f"Processing segment {s.id} with start time: {s.start} and end time: {s.end}")
+
+            # Process words in the segment
+            words = [{
+                'start': float(w.start),
+                'end': float(w.end),
+                'word': w.word,
+                'probability': float(w.probability)
+            } for w in s.words]
+
+            seg = {
+                'id': int(s.id),
+                'seek': int(s.seek),
+                'start': float(s.start),
+                'end': float(s.end),
+                'text': s.text,
+                'avg_logprob': float(s.avg_logprob),
+                'compression_ratio': float(s.compression_ratio),
+                'no_speech_prob': float(s.no_speech_prob),
+                'words': words
+            }
+            logger.debug(f'Adding new transcription segment: {seg}')
+            ret['segments'].append(seg)
+
+            logger.debug(f"Total segments in transcription result: {len(ret['segments'])}")
+            return ret
     except Exception as e:
-        logging.error(f"Error during transcription: {e}")
-        raise e
+        logger.error(f"Transcription error: {e}")
+        return {}
 
-    # Track the new segments and update the last transcribed time
-    for s in segs:
-        logging.info(f"Processing segment with start time: {s.start} and end time: {s.end}")
+def process_received_audio(data):
+    """Convert received bytes into normalized float32 NumPy array."""
+    logger.debug(f"Processing received audio data of size {len(data)} bytes")
+    audio_int16 = np.frombuffer(data, dtype=np.int16)
+    logger.debug(f"Converted to int16 NumPy array with {len(audio_int16)} samples")
 
-        # Only process segments that start after the last transcribed time
-        logging.info(f"New segment found starting at {s.start} seconds.")
-        words = [{'start': w.start, 'end': w.end, 'word': w.word, 'probability': w.probability} for w in s.words]
+    audio_float32 = audio_int16.astype(np.float32) / 32768.0  # Normalize to [-1, 1]
+    logger.debug(f"Normalized audio data to float32 with {len(audio_float32)} samples")
 
-        seg = {
-            'id': s.id, 'seek': s.seek, 'start': s.start, 'end': s.end, 'text': s.text,
-            'avg_logprob': s.avg_logprob, 'compression_ratio': s.compression_ratio,
-            'no_speech_prob': s.no_speech_prob, 'words': words
-        }
-        logging.info(f'Adding new transcription segment: {seg}')
-        ret['segements'].append(seg)
+    return audio_float32
 
-        # Update the last transcribed time to the end of the current segment
 
 
-#logging.info(f"Returning {len(ret['new_segments'])} new segments and updated last transcribed time.")
-    return ret
 
 
-import tempfile
 
 
-@app.websocket("/wtranscribe")
-async def websocket_transcribe(websocket: WebSocket):
-    logging.info("New WebSocket connection request received.")
-    await websocket.accept()
-    logging.info("WebSocket connection established successfully.")
 
-    try:
-        processed_segments = []  # Keeps track of the segments already transcribed
-        accumulated_audio_size = 0  # Track how much audio data has been buffered
-        accumulated_audio_time = 0  # Track the total audio duration accumulated
-        last_transcribed_time = 0.0
-        #min_transcription_time = 5.0  # Minimum duration of audio in seconds before transcription starts
-
-        # A temporary file to store the growing audio data
-
-        while True:
-            try:
-                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file: ##new temp file for every chunk
-                    logging.info(f"Temporary audio file created at {temp_audio_file.name}")
-                    # Receive the next chunk of audio data
-                    audio_chunk = await websocket.receive_bytes()
-                    if not audio_chunk:
-                        logging.warning("Received empty audio chunk, skipping processing.")
-                        continue
-
-                    # Write audio chunk to file and accumulate size and time
-                    temp_audio_file.write(audio_chunk)
-                    temp_audio_file.flush()
-                    accumulated_audio_size += len(audio_chunk)
-
-                    # Estimate the duration of the chunk based on its size (e.g., 16kHz audio)
-                    chunk_duration = len(audio_chunk) / (16000 * 2)  # Assuming 16kHz mono WAV (2 bytes per sample)
-                    accumulated_audio_time += chunk_duration
-
-                    partial_result = await transcribe_core_ws(temp_audio_file)
-                    accumulated_audio_time = 0  # Reset the accumulated audio time
-                    await websocket.send_json(partial_result)
-
-            except WebSocketDisconnect:
-                logging.info("WebSocket connection closed by the client.")
-                break
 
-    except Exception as e:
-        logging.error(f"Unexpected error during WebSocket transcription: {e}")
-        await websocket.send_json({"error": str(e)})
 
-    finally:
-        logging.info("Cleaning up and closing WebSocket connection.")
+# @app.websocket("/wtranscribe")
+# async def websocket_transcribe(websocket: WebSocket):
+#     logging.info("New WebSocket connection request received.")
+#     await websocket.accept()
+#     logging.info("WebSocket connection established successfully.")
+#
+#     try:
+#         while True:
+#             try:
+#                 audio_chunk = await websocket.receive_bytes()
+#                 if not audio_chunk:
+#                     logging.warning("Received empty audio chunk, skipping processing.")
+#                     continue
+#                 with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file: ##new temp file for every chunk
+#                     logging.info(f"Temporary audio file created at {temp_audio_file.name}")
+#                     # Receive the next chunk of audio data
+#
+#
+#
+#                     partial_result = await transcribe_core_ws(temp_audio_file.name)
+#                     await websocket.send_json(partial_result)
+#
+#             except WebSocketDisconnect:
+#                 logging.info("WebSocket connection closed by the client.")
+#                 break
+#
+#     except Exception as e:
+#         logging.error(f"Unexpected error during WebSocket transcription: {e}")
+#         await websocket.send_json({"error": str(e)})
+#
+#     finally:
+#         logging.info("Cleaning up and closing WebSocket connection.")