audio.py

import asyncio
import numpy as np
import ffmpeg
from time import time, sleep


from whisper_streaming_custom.whisper_online import online_factory
import math
import logging
import traceback
from state import SharedState
from formatters import format_time


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logging.getLogger().setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

class AudioProcessor:
    
    def __init__(self, args, asr, tokenizer):
        self.args = args
        self.sample_rate = 16000
        self.channels = 1
        self.samples_per_sec = int(self.sample_rate * args.min_chunk_size)
        self.bytes_per_sample = 2
        self.bytes_per_sec = self.samples_per_sec * self.bytes_per_sample
        self.max_bytes_per_sec = 32000 * 5  # 5 seconds of audio at 32 kHz
        
        
        self.shared_state = SharedState()
        self.asr = asr
        self.tokenizer = tokenizer
        
        self.ffmpeg_process = self.start_ffmpeg_decoder()

        self.transcription_queue = asyncio.Queue() if self.args.transcription else None
        self.diarization_queue = asyncio.Queue() if self.args.diarization else None

        self.pcm_buffer = bytearray()
        if self.args.transcription:
            self.online = online_factory(self.args, self.asr, self.tokenizer)

        

    def convert_pcm_to_float(self, pcm_buffer):
        """
        Converts a PCM buffer in s16le format to a normalized NumPy array.
        Arg: pcm_buffer. PCM buffer containing raw audio data in s16le format
        Returns: np.ndarray. NumPy array of float32 type normalized between -1.0 and 1.0
        """
        pcm_array = (np.frombuffer(pcm_buffer, dtype=np.int16).astype(np.float32) 
                    / 32768.0)
        return pcm_array

    async def start_ffmpeg_decoder(self):
        """
        Start an FFmpeg process in async streaming mode that reads WebM from stdin
        and outputs raw s16le PCM on stdout. Returns the process object.
        """
        process = (
            ffmpeg.input("pipe:0", format="webm")
            .output(
                "pipe:1",
                format="s16le",
                acodec="pcm_s16le",
                ac=self.channels,
                ar=str(self.sample_rate),
            )
            .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True)
        )
        return process

    async def restart_ffmpeg(self):
        if self.ffmpeg_process:
            try:
                self.ffmpeg_process.kill()
                await asyncio.get_event_loop().run_in_executor(None, self.ffmpeg_process.wait)
            except Exception as e:
                logger.warning(f"Error killing FFmpeg process: {e}")
            self.ffmpeg_process = await self.start_ffmpeg_decoder()
            self.pcm_buffer = bytearray()

    async def ffmpeg_stdout_reader(self):
        loop = asyncio.get_event_loop()
        beg = time()
        
        while True:
            try:
                elapsed_time = math.floor((time() - beg) * 10) / 10 # Round to 0.1 sec
                ffmpeg_buffer_from_duration = max(int(32000 * elapsed_time), 4096)
                beg = time()

                # Read chunk with timeout
                try:
                    chunk = await asyncio.wait_for(
                        loop.run_in_executor(
                            None, self.ffmpeg_process.stdout.read, ffmpeg_buffer_from_duration
                        ),
                        timeout=15.0
                    )
                except asyncio.TimeoutError:
                    logger.warning("FFmpeg read timeout. Restarting...")
                    await self.restart_ffmpeg()
                    beg = time()
                    continue  # Skip processing and read from new process

                if not chunk:
                    logger.info("FFmpeg stdout closed.")
                    break
                self.pcm_buffer.extend(chunk)
                        
                if self.args.diarization and self.diarization_queue:
                    await self.diarization_queue.put(self.convert_pcm_to_float(self.pcm_buffer).copy())

                if len(self.pcm_buffer) >= self.bytes_per_sec:
                    if len(self.pcm_buffer) > self.max_bytes_per_sec:
                        logger.warning(
                            f"""Audio buffer is too large: {len(self.pcm_buffer) / self.bytes_per_sec:.2f} seconds.
                            The model probably struggles to keep up. Consider using a smaller model.
                            """)

                    pcm_array = self.convert_pcm_to_float(self.pcm_buffer[:self.max_bytes_per_sec])
                    self.pcm_buffer = self.pcm_buffer[self.max_bytes_per_sec:]
                    
                    if self.args.transcription and self.transcription_queue:
                        await self.transcription_queue.put(pcm_array.copy())
                    
                    
                    if not self.args.transcription and not self.args.diarization:
                        await asyncio.sleep(0.1)
                    
            except Exception as e:
                logger.warning(f"Exception in ffmpeg_stdout_reader: {e}")
                logger.warning(f"Traceback: {traceback.format_exc()}")
                break
        logger.info("Exiting ffmpeg_stdout_reader...")

    async def transcription_processor(self):
        full_transcription = ""
        sep = self.online.asr.sep
        
        while True:
            try:
                pcm_array = await self.transcription_queue.get()
                
                logger.info(f"{len(self.online.audio_buffer) / self.online.SAMPLING_RATE} seconds of audio will be processed by the model.")
                
                # Process transcription
                self.online.insert_audio_chunk(pcm_array)
                new_tokens = self.online.process_iter()
                
                if new_tokens:
                    full_transcription += sep.join([t.text for t in new_tokens])
                    
                _buffer = self.online.get_buffer()
                buffer = _buffer.text
                end_buffer = _buffer.end if _buffer.end else (new_tokens[-1].end if new_tokens else 0)
                
                if buffer in full_transcription:
                    buffer = ""
                    
                await self.shared_state.update_transcription(
                    new_tokens, buffer, end_buffer, full_transcription, sep)
                
            except Exception as e:
                logger.warning(f"Exception in transcription_processor: {e}")
                logger.warning(f"Traceback: {traceback.format_exc()}")
            finally:
                self.transcription_queue.task_done()


    async def diarization_processor(self, diarization_obj):
        buffer_diarization = ""
        
        while True:
            try:
                pcm_array = await self.diarization_queue.get()
                
                # Process diarization
                await diarization_obj.diarize(pcm_array)
                
                # Get current state
                state = await self.shared_state.get_current_state()
                tokens = state["tokens"]
                end_attributed_speaker = state["end_attributed_speaker"]
                
                # Update speaker information
                new_end_attributed_speaker = diarization_obj.assign_speakers_to_tokens(
                    end_attributed_speaker, tokens)
                
                await self.shared_state.update_diarization(new_end_attributed_speaker, buffer_diarization)
                
            except Exception as e:
                logger.warning(f"Exception in diarization_processor: {e}")
                logger.warning(f"Traceback: {traceback.format_exc()}")
            finally:
                self.diarization_queue.task_done()

    async def results_formatter(self, websocket):
        while True:
            try:
                # Get the current state
                state = await self.shared_state.get_current_state()
                tokens = state["tokens"]
                buffer_transcription = state["buffer_transcription"]
                buffer_diarization = state["buffer_diarization"]
                end_attributed_speaker = state["end_attributed_speaker"]
                remaining_time_transcription = state["remaining_time_transcription"]
                remaining_time_diarization = state["remaining_time_diarization"]
                sep = state["sep"]
                
                # If diarization is enabled but no transcription, add dummy tokens periodically
                if (not tokens or tokens[-1].is_dummy) and not self.args.transcription and self.args.diarization:
                    await self.shared_state.add_dummy_token()
                    sleep(0.5)
                    state = await self.shared_state.get_current_state()
                    tokens = state["tokens"]
                # Process tokens to create response
                previous_speaker = -1
                lines = []
                last_end_diarized = 0
                undiarized_text = []
                
                for token in tokens:
                    speaker = token.speaker
                    if self.args.diarization:
                        if (speaker == -1 or speaker == 0) and token.end >= end_attributed_speaker:
                            undiarized_text.append(token.text)
                            continue
                        elif (speaker == -1 or speaker == 0) and token.end < end_attributed_speaker:
                            speaker = previous_speaker
                        if speaker not in [-1, 0]:
                            last_end_diarized = max(token.end, last_end_diarized)

                    if speaker != previous_speaker or not lines:
                        lines.append(
                            {
                                "speaker": speaker,
                                "text": token.text,
                                "beg": format_time(token.start),
                                "end": format_time(token.end),
                                "diff": round(token.end - last_end_diarized, 2)
                            }
                        )
                        previous_speaker = speaker
                    elif token.text:  # Only append if text isn't empty
                        lines[-1]["text"] += sep + token.text
                        lines[-1]["end"] = format_time(token.end)
                        lines[-1]["diff"] = round(token.end - last_end_diarized, 2)
                
                if undiarized_text:
                    combined_buffer_diarization = sep.join(undiarized_text)
                    if buffer_transcription:
                        combined_buffer_diarization += sep
                    await self.shared_state.update_diarization(end_attributed_speaker, combined_buffer_diarization)
                    buffer_diarization = combined_buffer_diarization
                    
                if lines:
                    response = {
                        "lines": lines, 
                        "buffer_transcription": buffer_transcription,
                        "buffer_diarization": buffer_diarization,
                        "remaining_time_transcription": remaining_time_transcription,
                        "remaining_time_diarization": remaining_time_diarization
                    }
                else:
                    response = {
                        "lines": [{
                            "speaker": 1,
                            "text": "",
                            "beg": format_time(0),
                            "end": format_time(tokens[-1].end) if tokens else format_time(0),
                            "diff": 0
                    }],
                        "buffer_transcription": buffer_transcription,
                        "buffer_diarization": buffer_diarization,
                        "remaining_time_transcription": remaining_time_transcription,
                        "remaining_time_diarization": remaining_time_diarization

                    }
                
                response_content = ' '.join([str(line['speaker']) + ' ' + line["text"] for line in lines]) + ' | ' + buffer_transcription + ' | ' + buffer_diarization
                
                if response_content != self.shared_state.last_response_content:
                    if lines or buffer_transcription or buffer_diarization:
                        await websocket.send_json(response)
                        self.shared_state.last_response_content = response_content
                
                # Add a small delay to avoid overwhelming the client
                await asyncio.sleep(0.1)
                
            except Exception as e:
                logger.warning(f"Exception in results_formatter: {e}")
                logger.warning(f"Traceback: {traceback.format_exc()}")
                await asyncio.sleep(0.5)  # Back off on error
                
    async def create_tasks(self, websocket, diarization):
        tasks = []    
        if self.args.transcription and self.online:
            tasks.append(asyncio.create_task(self.transcription_processor()))    
        if self.args.diarization and diarization:
            tasks.append(asyncio.create_task(self.diarization_processor(diarization)))
        formatter_task = asyncio.create_task(self.results_formatter(websocket))
        tasks.append(formatter_task)
        stdout_reader_task = asyncio.create_task(self.ffmpeg_stdout_reader())
        tasks.append(stdout_reader_task)
        self.tasks = tasks
        self.diarization = diarization
        
    async def cleanup(self):
        for task in self.tasks:
            task.cancel()
        try:
            await asyncio.gather(*self.tasks, return_exceptions=True)
            self.ffmpeg_process.stdin.close()
            self.ffmpeg_process.wait()
        except Exception as e:
            logger.warning(f"Error during cleanup: {e}")
        if self.args.diarization and self.diarization:
            self.diarization.close()

    async def process_audio(self, message):
        try:
            self.ffmpeg_process.stdin.write(message)
            self.ffmpeg_process.stdin.flush()
        except (BrokenPipeError, AttributeError) as e:
            logger.warning(f"Error writing to FFmpeg: {e}. Restarting...")
            await self.restart_ffmpeg()
            self.ffmpeg_process.stdin.write(message)
            self.ffmpeg_process.stdin.flush()