app.py

# server.py remains the same as before

# Updated client.py
import asyncio
import websockets
import numpy as np
import base64
import argparse
import requests
import time
import torch
import torchaudio

import av
import streamlit as st
from typing import List
from streamlit_webrtc import WebRtcMode, webrtc_streamer

class AudioClient:
    def __init__(self, server_url="ws://localhost:8000", token_temp=None, categorical_temp=None, gaussian_temp=None):
        # Convert ws:// to http:// for the base URL
        self.base_url = server_url.replace("ws://", "http://")
        self.server_url = f"{server_url}/audio"
        self.sound_check = False
        
        # Set temperatures if provided
        if any(t is not None for t in [token_temp, categorical_temp, gaussian_temp]):
            response_message = self.set_temperature_and_echo(token_temp, categorical_temp, gaussian_temp)
            print(response_message)

        self.downsampler = torchaudio.transforms.Resample(STREAMING_SAMPLE_RATE, SAMPLE_RATE)
        self.upsampler = torchaudio.transforms.Resample(SAMPLE_RATE, STREAMING_SAMPLE_RATE)
        self.ws = None
        self.in_buffer = None
        self.out_buffer = None
    
    def set_temperature_and_echo(self, token_temp=None, categorical_temp=None, gaussian_temp=None, echo_testing = False):
        """Send temperature settings to server"""
        params = {}
        if token_temp is not None:
            params['token_temp'] = token_temp
        if categorical_temp is not None:
            params['categorical_temp'] = categorical_temp
        if gaussian_temp is not None:
            params['gaussian_temp'] = gaussian_temp
            
        response = requests.post(f"{self.base_url}/set_temperature", params=params)
        response_message = response.json()['message']
        return response_message
    
    def _resample(self, audio_data: np.ndarray, resampler: torchaudio.transforms.Resample) -> np.ndarray:
        audio_data = audio_data.astype(np.float32) / 32767.0
        audio_data = resampler(torch.tensor(audio_data)).numpy()
        audio_data = (audio_data * 32767.0).astype(np.int16)
        return audio_data
    
    def upsample(self, audio_data: np.ndarray) -> np.ndarray:
        return self._resample(audio_data, self.upsampler)
    
    def downsample(self, audio_data: np.ndarray) -> np.ndarray:
        return self._resample(audio_data, self.downsampler)
    
    def from_s16_format(self, audio_data: np.ndarray, channels: int) -> np.ndarray:
        if channels == 2:
            audio_data = audio_data.reshape(-1, 2).T
        else:
            audio_data = audio_data.reshape(-1)
        return audio_data
    
    def to_s16_format(self, audio_data: np.ndarray):
        if len(audio_data.shape) == 2 and audio_data.shape[0] == 2:
            audio_data = audio_data.T.reshape(1, -1)
        elif len(audio_data.shape) == 1:
            audio_data = audio_data.reshape(1, -1)
        return audio_data
    
    def to_channels(self, audio_data: np.ndarray, channels: int) -> np.ndarray:
        current_channels = audio_data.shape[0] if len(audio_data.shape) == 2 else 1
        if current_channels == channels:
            return audio_data
        elif current_channels == 1 and channels == 2:
            audio_data = np.tile(audio_data, 2).reshape(2, -1)
        elif current_channels == 2 and channels == 1:
            audio_data = audio_data.astype(np.float32) / 32767.0
            audio_data = audio_data.mean(axis=0)
            audio_data = (audio_data * 32767.0).astype(np.int16)
        return audio_data

    async def process_audio(self, audio_data: np.ndarray) -> np.ndarray:
        if self.ws is None:
            self.ws = await websockets.connect(self.server_url)

        audio_data = audio_data.reshape(-1, CHANNELS)
        print(f'Data from microphone:{audio_data.shape, audio_data.dtype, audio_data.min(), audio_data.max()}')

        # Convert to base64
        audio_b64 = base64.b64encode(audio_data.tobytes()).decode('utf-8')
        
        # Send to server
        time_sent = time.time()
        await self.ws.send(f"data:audio/raw;base64,{audio_b64}")
    
        # Receive processed audio
        response = await self.ws.recv()
        response = response.split(",")[1]
        time_received = time.time()
        print(f"Data sent: {audio_b64[:10]}. Data received: {response[:10]}. Received in {(time_received - time_sent) * 1000:.2f} ms")
        processed_audio = np.frombuffer(
            base64.b64decode(response),
            dtype=np.int16
        ).reshape(-1, CHANNELS)
        print(f'Data from model:{processed_audio.shape, processed_audio.dtype, processed_audio.min(), processed_audio.max()}')

        if CHANNELS == 1:
            processed_audio = processed_audio.reshape(-1)
        return processed_audio
    
    async def queued_audio_frames_callback(self, frames: List[av.AudioFrame]) -> List[av.AudioFrame]:
        out_frames = []
        for frame in frames:
            # Read in audio
            audio_data = frame.to_ndarray()

            # Convert input audio from s16 format, convert to `CHANNELS` number of channels, and downsample
            audio_data = self.from_s16_format(audio_data, len(frame.layout.channels))
            audio_data = self.to_channels(audio_data, CHANNELS)
            audio_data = self.downsample(audio_data)

            # Add audio to input buffer
            if self.in_buffer is None:
                self.in_buffer = audio_data
            else:
                self.in_buffer = np.concatenate((self.in_buffer, audio_data), axis=-1)
            
            # Take BLOCK_SIZE samples from input buffer if available for processing
            if self.in_buffer.shape[0] >= BLOCK_SIZE:
                audio_data = self.in_buffer[:BLOCK_SIZE]
                self.in_buffer = self.in_buffer[BLOCK_SIZE:]
            else:
                audio_data = None
            
            # Process audio if available and add resulting audio to output buffer
            if audio_data is not None:
                if not self.sound_check:
                    audio_data = await self.process_audio(audio_data)
                if self.out_buffer is None:
                    self.out_buffer = audio_data
                else:
                    self.out_buffer = np.concatenate((self.out_buffer, audio_data), axis=-1)

            # Take `out_samples` samples from output buffer if available for output
            out_samples = int(frame.samples * SAMPLE_RATE / STREAMING_SAMPLE_RATE)
            if self.out_buffer is not None and self.out_buffer.shape[0] >= out_samples:
                audio_data = self.out_buffer[:out_samples]
                self.out_buffer = self.out_buffer[out_samples:]
            else:
                audio_data = None

            # Output silence if no audio data available
            if audio_data is None:
                # output silence
                audio_data = np.zeros(out_samples, dtype=np.int16)
            
            # Upsample output audio, convert to original number of channels, and convert to s16 format
            audio_data = self.upsample(audio_data)
            audio_data = self.to_channels(audio_data, len(frame.layout.channels))
            audio_data = self.to_s16_format(audio_data)

            # return audio data as AudioFrame
            new_frame = av.AudioFrame.from_ndarray(audio_data, format=frame.format.name, layout=frame.layout.name)
            new_frame.sample_rate = frame.sample_rate
            out_frames.append(new_frame)

        return out_frames

    def stop(self): 
        if self.ws is not None:
            # TODO: this hangs. Figure out why.
            #asyncio.get_event_loop().run_until_complete(self.ws.close())
            print("Websocket closed")
        self.ws = None
        self.in_buffer = None
        self.out_buffer = None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Audio Client with Temperature Control')
    parser.add_argument('--token_temp', '-t1', type=float, help='Token (LM) temperature parameter')
    parser.add_argument('--categorical_temp', '-t2', type=float, help='Categorical (VAE) temperature parameter')
    parser.add_argument('--gaussian_temp', '-t3', type=float, help='Gaussian (VAE) temperature parameter')
    parser.add_argument('--server', '-s', default="ws://localhost:8000", 
                        help='Server URL (default: ws://localhost:8000)')
    parser.add_argument("--use_ice_servers", action="store_true", help="Use public STUN servers")
    
    args = parser.parse_args()
    
    # Audio settings
    STREAMING_SAMPLE_RATE = 48000
    SAMPLE_RATE = 16000
    BLOCK_SIZE = 2000
    CHANNELS = 1
    
    st.title("hertz-dev webrtc demo!")
    st.markdown("""
    Welcome to the audio processing interface! Here you can talk live with hertz.
    - Process audio in real-time through your microphone
    - Adjust various temperature parameters for inference
    - Test your microphone with sound check mode
    - Enable/disable echo cancellation and noise suppression
    
    To begin, click the START button below and allow microphone access.
    """)

    audio_client = st.session_state.get("audio_client")
    if audio_client is None:
        audio_client = AudioClient(
            server_url=args.server,
            token_temp=args.token_temp,
            categorical_temp=args.categorical_temp,
            gaussian_temp=args.gaussian_temp
        )
        st.session_state.audio_client = audio_client

    with st.sidebar:
        st.markdown("## Inference Settings")
        token_temp_default = args.token_temp if args.token_temp is not None else 0.8
        token_temp = st.slider("Token Temperature", 0.05, 2.0, token_temp_default, step=0.05)
        categorical_temp_default = args.categorical_temp if args.categorical_temp is not None else 0.4
        categorical_temp = st.slider("Categorical Temperature", 0.01, 1.0, categorical_temp_default, step=0.01)
        gaussian_temp_default = args.gaussian_temp if args.gaussian_temp is not None else 0.1
        gaussian_temp = st.slider("Gaussian Temperature", 0.01, 1.0, gaussian_temp_default, step=0.01)
        if st.button("Set Temperatures"):
            response_message = audio_client.set_temperature_and_echo(token_temp, categorical_temp, gaussian_temp)
            st.write(response_message)

        st.markdown("## Microphone Settings")
        audio_client.sound_check = st.toggle("Sound Check (Echo)", value=False)
        echo_cancellation = st.toggle("Echo Cancellation*‡", value=False)
        noise_suppression = st.toggle("Noise Suppression*", value=False)
        st.markdown(r"\* *Restart stream to take effect*")
        st.markdown("‡ *May cause audio to cut out*")

    # Use a free STUN server from Google if --use_ice_servers is given
    # (found in get_ice_servers() at https://github.com/whitphx/streamlit-webrtc/blob/main/sample_utils/turn.py)
    rtc_configuration = {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]} if args.use_ice_servers else None
    audio_config = {"echoCancellation": echo_cancellation, "noiseSuppression": noise_suppression}
    webrtc_streamer(
        key="streamer",
        mode=WebRtcMode.SENDRECV,
        rtc_configuration=rtc_configuration,
        media_stream_constraints={"audio": audio_config, "video": False},
        queued_audio_frames_callback=audio_client.queued_audio_frames_callback,
        on_audio_ended=audio_client.stop,
        async_processing=True,
    )