fix: make update_audio fn async to avoid server freeze

src/audio/audio_generator.py

@@ -1,18 +1,14 @@
 import asyncio
 from google.genai import types
 import wave
-import queue
 import logging
 import io
-import time
+import gradio as gr
 from config import settings
 from services.google import GoogleClientFactory
 
 logger = logging.getLogger(__name__)
 
-
-
-
 async def generate_music(user_hash: str, music_tone: str, receive_audio):
     if user_hash in sessions:
         logger.info(
@@ -44,7 +40,7 @@ async def generate_music(user_hash: str, music_tone: str, receive_audio):
             logger.info(
                 f"Started music generation for user hash {user_hash}, music tone: {music_tone}"
             )
-            sessions[user_hash] = {"session": session, "queue": queue.Queue()}
+            sessions[user_hash] = {"session": session, "queue": asyncio.Queue()}
 
 
 async def change_music_tone(user_hash: str, new_tone):
@@ -75,7 +71,7 @@ async def receive_audio(session, user_hash):
                     audio_data = message.server_content.audio_chunks[0].data
                     queue = sessions[user_hash]["queue"]
                     # audio_data is already bytes (raw PCM)
-                    await asyncio.to_thread(queue.put, audio_data)
+                    await queue.put(audio_data)
                 await asyncio.sleep(10**-12)
         except Exception as e:
             logger.error(f"Error in receive_audio: {e}")
@@ -102,44 +98,66 @@ async def cleanup_music_session(user_hash: str):
         del sessions[user_hash]
 
 
-def update_audio(user_hash):
-    """Continuously stream audio from the queue as WAV bytes."""
+async def update_audio(user_hash: str, request: gr.Request):
+    """
+    Continuously stream audio from the queue as WAV bytes, and clean up
+    when the user disconnects.
+    """
     if user_hash == "":
         return
 
     logger.info(f"Starting audio update loop for user hash: {user_hash}")
-    while True:
-        if user_hash not in sessions:
-            time.sleep(0.5)
-            continue
-        queue = sessions[user_hash]["queue"]
-        pcm_data = queue.get()  # This is raw PCM audio bytes
-
-        if not isinstance(pcm_data, bytes):
-            logger.warning(
-                f"Expected bytes from audio_queue, got {type(pcm_data)}. Skipping."
-            )
-            continue
-
-        # Lyria provides stereo, 16-bit PCM at 48kHz.
-        # Ensure the number of bytes is consistent with stereo 16-bit audio.
-        # Each frame = NUM_CHANNELS * SAMPLE_WIDTH bytes.
-        # If len(pcm_data) is not a multiple of (NUM_CHANNELS * SAMPLE_WIDTH),
-        # it might indicate an incomplete chunk or an issue.
-        bytes_per_frame = NUM_CHANNELS * SAMPLE_WIDTH
-        if len(pcm_data) % bytes_per_frame != 0:
-            logger.warning(
-                f"Received PCM data with length {len(pcm_data)}, which is not a multiple of "
-                f"bytes_per_frame ({bytes_per_frame}). This might cause issues with WAV formatting."
-            )
-            # Depending on strictness, you might want to skip this chunk:
-            # continue
-
-        wav_buffer = io.BytesIO()
-        with wave.open(wav_buffer, "wb") as wf:
-            wf.setnchannels(NUM_CHANNELS)
-            wf.setsampwidth(SAMPLE_WIDTH)  # Corresponds to 16-bit audio
-            wf.setframerate(SAMPLE_RATE)
-            wf.writeframes(pcm_data)
-        wav_bytes = wav_buffer.getvalue()
-        yield wav_bytes
+    try:
+        while True:
+            if await request.request.is_disconnected():
+                logger.info(f"Client disconnected for user hash {user_hash}.")
+                break
+
+            if user_hash not in sessions:
+                await asyncio.sleep(0.5)
+                continue
+
+            try:
+                queue = sessions[user_hash]["queue"]
+                pcm_data = await asyncio.wait_for(queue.get(), timeout=1.0)
+            except asyncio.TimeoutError:
+                continue  # Check for disconnect again
+            except (KeyError, AttributeError):
+                logger.warning(
+                    f"Session or queue for {user_hash} not found. Stopping audio loop."
+                )
+                break
+
+            if not isinstance(pcm_data, bytes):
+                logger.warning(
+                    f"Expected bytes from audio_queue, got {type(pcm_data)}. Skipping."
+                )
+                continue
+
+            # Lyria provides stereo, 16-bit PCM at 48kHz.
+            # Ensure the number of bytes is consistent with stereo 16-bit audio.
+            # Each frame = NUM_CHANNELS * SAMPLE_WIDTH bytes.
+            # If len(pcm_data) is not a multiple of (NUM_CHANNELS * SAMPLE_WIDTH),
+            # it might indicate an incomplete chunk or an issue.
+            bytes_per_frame = NUM_CHANNELS * SAMPLE_WIDTH
+            if len(pcm_data) % bytes_per_frame != 0:
+                logger.warning(
+                    f"Received PCM data with length {len(pcm_data)}, which is not a multiple of "
+                    f"bytes_per_frame ({bytes_per_frame}). This might cause issues with WAV formatting."
+                )
+                # Depending on strictness, you might want to skip this chunk:
+                # continue
+
+            wav_buffer = io.BytesIO()
+            with wave.open(wav_buffer, "wb") as wf:
+                wf.setnchannels(NUM_CHANNELS)
+                wf.setsampwidth(SAMPLE_WIDTH)  # Corresponds to 16-bit audio
+                wf.setframerate(SAMPLE_RATE)
+                wf.writeframes(pcm_data)
+            wav_bytes = wav_buffer.getvalue()
+            yield wav_bytes
+    finally:
+        logger.info(
+            f"Audio update loop finished for {user_hash}. Cleaning up music session."
+        )
+        await cleanup_music_session(user_hash)

src/main.py

@@ -357,7 +357,6 @@ with gr.Blocks(
         outputs=[game_text, game_image, game_choices, custom_choice],
     )
 
-    demo.unload(cleanup_music_session)
     demo.load(
         fn=generate_user_hash,
         inputs=[],