Qwen-Omni-Realtime

# -- coding: utf-8 --
import os
import asyncio
import time
import threading
import queue
import pyaudio
from omni_realtime_client import OmniRealtimeClient, TurnDetectionMode

class AudioPlayer:
    """Real-time audio player class."""

    def __init__(self, sample_rate=24000, channels=1, sample_width=2):
        self.sample_rate = sample_rate
        self.channels = channels
        self.sample_width = sample_width  # 2 bytes for 16-bit
        self.audio_queue = queue.Queue()
        self.is_playing = False
        self.play_thread = None
        self.pyaudio_instance = None
        self.stream = None
        self._lock = threading.Lock()  # Add a lock for synchronized access.
        self._last_data_time = time.time()  # Record the time when the last data was received.
        self._response_done = False  # Add a flag to indicate that the response is complete.
        self._waiting_for_response = False  # Flag to indicate whether the client is waiting for a server response.
        # Record the time of the last write to the audio stream and the duration of the last audio chunk for more accurate playback end detection.
        self._last_play_time = time.time()
        self._last_chunk_duration = 0.0

    def start(self):
        """Start the audio player."""
        with self._lock:
            if self.is_playing:
                return

            self.is_playing = True

            try:
                self.pyaudio_instance = pyaudio.PyAudio()

                # Create an audio output stream.
                self.stream = self.pyaudio_instance.open(
                    format=pyaudio.paInt16,  # 16-bit
                    channels=self.channels,
                    rate=self.sample_rate,
                    output=True,
                    frames_per_buffer=1024
                )

                # Start the playback thread.
                self.play_thread = threading.Thread(target=self._play_audio)
                self.play_thread.daemon = True
                self.play_thread.start()

                print("Audio player started")
            except Exception as e:
                print(f"Failed to start audio player: {e}")
                self._cleanup_resources()
                raise

    def stop(self):
        """Stop the audio player."""
        with self._lock:
            if not self.is_playing:
                return

            self.is_playing = False

        # Clear the queue.
        while not self.audio_queue.empty():
            try:
                self.audio_queue.get_nowait()
            except queue.Empty:
                break

        # Wait for the playback thread to finish. Wait outside the lock to avoid a deadlock.
        if self.play_thread and self.play_thread.is_alive():
            self.play_thread.join(timeout=2.0)

        # Acquire the lock again to clean up resources.
        with self._lock:
            self._cleanup_resources()

        print("Audio player stopped")

    def _cleanup_resources(self):
        """Clean up audio resources. This must be called within the lock."""
        try:
            # Close the audio stream.
            if self.stream:
                if not self.stream.is_stopped():
                    self.stream.stop_stream()
                self.stream.close()
                self.stream = None
        except Exception as e:
            print(f"Error closing audio stream: {e}")

        try:
            if self.pyaudio_instance:
                self.pyaudio_instance.terminate()
                self.pyaudio_instance = None
        except Exception as e:
            print(f"Error terminating PyAudio: {e}")

    def add_audio_data(self, audio_data):
        """Add audio data to the playback queue."""
        if self.is_playing and audio_data:
            self.audio_queue.put(audio_data)
            with self._lock:
                self._last_data_time = time.time()  # Update the time when the last data was received.
                self._waiting_for_response = False  # Data received, no longer waiting.

    def stop_receiving_data(self):
        """Mark that no more new audio data will be received."""
        with self._lock:
            self._response_done = True
            self._waiting_for_response = False  # Response ended, no longer waiting.

    def prepare_for_next_turn(self):
        """Reset the player state for the next conversation turn."""
        with self._lock:
            self._response_done = False
            self._last_data_time = time.time()
            self._last_play_time = time.time()
            self._last_chunk_duration = 0.0
            self._waiting_for_response = True  # Start waiting for the next response.

        # Clear any remaining audio data from the previous turn.
        while not self.audio_queue.empty():
            try:
                self.audio_queue.get_nowait()
            except queue.Empty:
                break

    def is_finished_playing(self):
        """Check if all audio data has been played."""
        with self._lock:
            queue_size = self.audio_queue.qsize()
            time_since_last_data = time.time() - self._last_data_time
            time_since_last_play = time.time() - self._last_play_time

            # ---------------------- Smart end detection ----------------------
            # 1. Preferred method: If the server has marked completion and the playback queue is empty,
            #    wait for the most recent audio chunk to finish playing (chunk duration + 0.1s tolerance).
            if self._response_done and queue_size == 0:
                min_wait = max(self._last_chunk_duration + 0.1, 0.5)  # Wait at least 0.5s.
                if time_since_last_play >= min_wait:
                    return True

            # 2. Fallback method: If no new data has been received for over a second and the playback queue is empty.
            #    This logic serves as a safeguard if the server does not explicitly send `response.done`.
            if not self._waiting_for_response and queue_size == 0 and time_since_last_data > 1.0:
                print("\n(No new audio received for a while, assuming playback is finished)")
                return True

            return False

    def _play_audio(self):
        """Worker thread for playing audio data."""
        while True:
            # Check if it should stop.
            with self._lock:
                if not self.is_playing:
                    break
                stream_ref = self.stream  # Get a reference to the stream.

            try:
                # Get audio data from the queue, with a timeout of 0.1 seconds.
                audio_data = self.audio_queue.get(timeout=0.1)

                # Check the status and stream validity again.
                with self._lock:
                    if self.is_playing and stream_ref and not stream_ref.is_stopped():
                        try:
                            # Play the audio data.
                            stream_ref.write(audio_data)
                            # Update the latest playback information.
                            self._last_play_time = time.time()
                            self._last_chunk_duration = len(audio_data) / (
                                        self.channels * self.sample_width) / self.sample_rate
                        except Exception as e:
                            print(f"Error writing to audio stream: {e}")
                            break

                # Mark this data block as processed.
                self.audio_queue.task_done()

            except queue.Empty:
                # Continue waiting if the queue is empty.
                continue
            except Exception as e:
                print(f"Error playing audio: {e}")
                break

class MicrophoneRecorder:
    """Real-time microphone recorder."""

    def __init__(self, sample_rate=16000, channels=1, chunk_size=3200):
        self.sample_rate = sample_rate
        self.channels = channels
        self.chunk_size = chunk_size
        self.pyaudio_instance = None
        self.stream = None
        self.frames = []
        self._is_recording = False
        self._record_thread = None

    def _recording_thread(self):
        """Recording worker thread."""
        # Continuously read data from the audio stream while _is_recording is True.
        while self._is_recording:
            try:
                # Use exception_on_overflow=False to avoid crashing due to a buffer overflow.
                data = self.stream.read(self.chunk_size, exception_on_overflow=False)
                self.frames.append(data)
            except (IOError, OSError) as e:
                # When the stream is closed, a read operation may raise an error.
                print(f"Error reading from recording stream, it might be closed: {e}")
                break

    def start(self):
        """Start recording."""
        if self._is_recording:
            print("Recording is already in progress.")
            return

        self.frames = []
        self._is_recording = True

        try:
            self.pyaudio_instance = pyaudio.PyAudio()
            self.stream = self.pyaudio_instance.open(
                format=pyaudio.paInt16,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                frames_per_buffer=self.chunk_size
            )

            self._record_thread = threading.Thread(target=self._recording_thread)
            self._record_thread.daemon = True
            self._record_thread.start()
            print("Microphone recording started...")
        except Exception as e:
            print(f"Failed to start microphone: {e}")
            self._is_recording = False
            self._cleanup()
            raise

    def stop(self):
        """Stop recording and return the audio data."""
        if not self._is_recording:
            return None

        self._is_recording = False

        # Wait for the recording thread to exit safely.
        if self._record_thread:
            self._record_thread.join(timeout=1.0)

        self._cleanup()

        print("Microphone recording stopped.")
        return b''.join(self.frames)

    def _cleanup(self):
        """Safely clean up PyAudio resources."""
        if self.stream:
            try:
                if self.stream.is_active():
                    self.stream.stop_stream()
                self.stream.close()
            except Exception as e:
                print(f"Error closing audio stream: {e}")

        if self.pyaudio_instance:
            try:
                self.pyaudio_instance.terminate()
            except Exception as e:
                print(f"Error terminating PyAudio instance: {e}")

        self.stream = None
        self.pyaudio_instance = None

async def interactive_test():
    """
    Interactive test for multi-turn conversations with audio and image support.
    """
    # ------------------- 1. Initialization and connection (one-time) -------------------
    # API keys for the Singapore and China (Beijing) regions are different. To get an API key, see https://www.alibabacloud.com/help/model-studio/get-api-key
    api_key = os.environ.get("DASHSCOPE_API_KEY")
    if not api_key:
        print("Please set the DASHSCOPE_API_KEY environment variable.")
        return

    print("--- Real-time Multimodal Audio/Video Chat Client ---")
    print("Initializing audio player and client...")

    audio_player = AudioPlayer()
    audio_player.start()

    def on_audio_received(audio_data):
        audio_player.add_audio_data(audio_data)

    def on_response_done(event):
        print("\n(Received response end marker)")
        audio_player.stop_receiving_data()

    realtime_client = OmniRealtimeClient(
        # This is the base_url for the Singapore region. If you use a model in the China (Beijing) region, replace the base_url with wss://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/api-ws/v1/realtime.
        base_url="wss://{WorkspaceId}.ap-southeast-1.maas.aliyuncs.com/api-ws/v1/realtime",
        api_key=api_key,
        model="qwen3.8-omni-flash-realtime",
        voice="Tina",
        instructions="You are Xiaoyun, a personal assistant. Please answer the user's questions accurately and in a friendly manner, always responding with a helpful attitude.", # Set the model's role.
        on_text_delta=lambda text: print(f"assistant: {text}", end="", flush=True),
        on_audio_delta=on_audio_received,
        turn_detection_mode=TurnDetectionMode.MANUAL,
        extra_event_handlers={"response.done": on_response_done}
    )

    message_handler_task = None
    try:
        await realtime_client.connect()
        print("Connected to the server. Enter 'q' or 'quit' to exit at any time.")
        message_handler_task = asyncio.create_task(realtime_client.handle_messages())
        await asyncio.sleep(0.5)

        turn_counter = 1
        # ------------------- 2. Multi-turn conversation loop -------------------
        while True:
            print(f"\n--- Turn {turn_counter} ---")
            audio_player.prepare_for_next_turn()

            recorded_audio = None
            image_paths = []

            # --- Get user input: Record from microphone ---
            loop = asyncio.get_event_loop()
            recorder = MicrophoneRecorder(sample_rate=16000)  # We recommend a 16k sample rate for speech recognition.

            print("Ready to record. Press Enter to start recording (or enter 'q' to exit)...")
            user_input = await loop.run_in_executor(None, input)
            if user_input.strip().lower() in ['q', 'quit']:
                print("User requested to exit...")
                return

            try:
                recorder.start()
            except Exception:
                print("Could not start recording. Please check your microphone permissions and device. Skipping this turn.")
                continue

            print("Recording... Press Enter again to stop recording.")
            await loop.run_in_executor(None, input)

            recorded_audio = recorder.stop()

            if not recorded_audio or len(recorded_audio) == 0:
                print("No valid audio was recorded. Please start this turn again.")
                continue

            # --- Get image input (optional) ---
            # The image input feature is disabled by default. Uncomment the code below to enable it.
            # print("\nEnter the absolute path of an [image file] on each line (optional). When finished, enter 's' or press Enter to send the request.")
            # while True:
            #     path = input("Image path: ").strip()
            #     if path.lower() == 's' or path == '':
            #         break
            #     if path.lower() in ['q', 'quit']:
            #         print("User requested to exit...")
            #         return
            #
            #     if not os.path.isabs(path):
            #         print("Error: Please enter an absolute path.")
            #         continue
            #     if not os.path.exists(path):
            #         print(f"Error: File not found -> {path}")
            #         continue
            #     image_paths.append(path)
            #     print(f"Image added: {os.path.basename(path)}")

            # --- 3. Send data and get response ---
            print("\n--- Input Confirmation ---")
            print(f"Audio to process: 1 (from microphone), Images: {len(image_paths)}")
            print("------------------")

            # 3.1 Send the recorded audio.
            try:
                print(f"Sending microphone recording ({len(recorded_audio)} bytes)")
                await realtime_client.stream_audio(recorded_audio)
                await asyncio.sleep(0.1)
            except Exception as e:
                print(f"Failed to send microphone recording: {e}")
                continue

            # 3.2 Send all image files.
            # The image input feature is disabled by default. Uncomment the code below to enable it.
            # for i, path in enumerate(image_paths):
            #     try:
            #         with open(path, "rb") as f:
            #             data = f.read()
            #         print(f"Sending image {i+1}: {os.path.basename(path)} ({len(data)} bytes)")
            #         await realtime_client.append_image(data)
            #         await asyncio.sleep(0.1)
            #     except Exception as e:
            #         print(f"Failed to send image {os.path.basename(path)}: {e}")

            # 3.3 Submit and wait for the response.
            print("Submitting all inputs, requesting server response...")
            await realtime_client.commit_audio_buffer()
            await realtime_client.create_response()

            print("Waiting for and playing server response audio...")
            start_time = time.time()
            max_wait_time = 60
            while not audio_player.is_finished_playing():
                if time.time() - start_time > max_wait_time:
                    print(f"\nWait timed out ({max_wait_time} seconds). Moving to the next turn.")
                    break
                await asyncio.sleep(0.2)

            print("\nAudio playback for this turn is complete!")
            turn_counter += 1

    except (asyncio.CancelledError, KeyboardInterrupt):
        print("\nProgram was interrupted.")
    except Exception as e:
        print(f"An unhandled error occurred: {e}")
    finally:
        # ------------------- 4. Clean up resources -------------------
        print("\nClosing connection and cleaning up resources...")
        if message_handler_task and not message_handler_task.done():
            message_handler_task.cancel()

        if 'realtime_client' in locals() and realtime_client.ws and not realtime_client.ws.close:
            await realtime_client.close()
            print("Connection closed.")

        audio_player.stop()
        print("Program exited.")

if __name__ == "__main__":
    try:
        asyncio.run(interactive_test())
    except KeyboardInterrupt:
        print("\nProgram was forcibly exited by the user.")