va_reader_omni.py

import datetime
import json
import os
import random
import subprocess
import threading
import time
import traceback
from collections import deque
from copy import deepcopy

import jsonschema
import numpy as np
import torch
import torch.distributed as dist
import zmq
from bson import BSON
from loguru import logger
from scipy.signal import resample


class AudioInfo:
    def __init__(self, info: dict):
        self.sample_count = info["sample_count"]
        self.sample_rate = info["sample_rate"]
        self.channel_count = info["channel_count"]
        self.sample_fmt = info["sample_fmt"]
        self.pts = info["pts"]

    def is_spec_equal(self, other: "AudioInfo") -> bool:
        return self.sample_fmt == other.sample_fmt and self.sample_rate == other.sample_rate and self.channel_count == other.channel_count

    def duration(self) -> datetime.timedelta:
        return datetime.timedelta(seconds=self.sample_count / self.sample_rate)

    def __str__(self):
        return "AudioInfo(sample_count={}, sample_rate={}, channel_count={}, sample_fmt={}, pts={})".format(self.sample_count, self.sample_rate, self.channel_count, self.sample_fmt, self.pts)


class ByteBuffer:
    def __init__(self):
        self.buffer = deque()
        self.current_size = 0
        # is the audio belonging to current turn finished
        self.audio_finished = False

    def add(self, byte_data: bytes):
        self.buffer.append(byte_data)
        self.current_size += len(byte_data)

    def get(self, size=1024):
        data = bytearray()

        while size > 0 and len(self.buffer) > 0:
            chunk = self.buffer.popleft()
            if len(chunk) <= size:
                # 如果当前数据小于size，则将当前数据全部添加到data中
                data.extend(chunk)
                self.current_size -= len(chunk)
                size -= len(chunk)
            else:
                # 如果当前数据大于size，则将当前数据的一部分添加到data中，剩余部分留在缓冲区
                data.extend(chunk[:size])
                self.buffer.appendleft(chunk[size:])  # 剩余部分留在缓冲区
                self.current_size -= size
                size = 0

        return bytes(data)

    def mark_finished(self):
        self.audio_finished = True

    def has_more_voice(self):
        return not self.audio_finished

    def __len__(self):
        return self.current_size


class ChatAdapter:
    def __init__(
        self,
        omni_work_dir: str,
        whep_url: str,
        session_id: str,
        account: str,
        config_files: list[str],
        config_schema_path: str,
        seg_duration: float,
        model_runner,
        huoshan_tts_voice_type,
    ):
        assert os.path.exists(omni_work_dir), f"OMNI work directory {omni_work_dir} does not exist"
        self.omni_work_dir = omni_work_dir
        self.context = zmq.Context()
        self.w2f_socket = self.context.socket(zmq.PULL)
        self.w2f_url = ChatAdapter.select_and_bind(self.w2f_socket)
        self.f2w_socket = self.context.socket(zmq.PUSH)
        self.f2w_url = ChatAdapter.select_and_bind(self.f2w_socket)
        self.recv_thread = None
        self.audio_buffer = ByteBuffer()
        self.audio_info = None
        self.chat_server_cmd = [
            os.path.join(self.omni_work_dir, "bin", "seko-chatter"),
            "--session-id",
            session_id,
            "--account",
            account,
            "--whep-server-url",
            whep_url,
            "--w2f-endpoint",
            self.w2f_url,
            "--f2w-endpoint",
            self.f2w_url,
            "--config-files",
            *config_files,
        ]
        override_config = {}
        if huoshan_tts_voice_type is not None:
            logger.info(f"Use Huoshan TTS voice type: {huoshan_tts_voice_type}")
            override_config["TTS"] = {
                "default_voice_info": {
                    "voice_type": huoshan_tts_voice_type,
                    "provider": "huoshan_stream_tts",
                }
            }
        with open(config_schema_path, "r") as f:
            schema = json.load(f)
        jsonschema.validate(instance=override_config, schema=schema)
        if override_config is not None:
            self.chat_server_cmd.extend(["--override-config", json.dumps(override_config)])
        self.chatter_proc = None

        self.seg_duration = seg_duration
        self.reset_prev = False
        self.status = "blank"
        self.immediate_switch = 0
        self.model_runner = model_runner

    def launch_chat_server(self):
        env = {
            "RUST_LOG": "info,duplex_server=debug,backend_5o=debug",
            "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", "") + ":" + os.path.join(self.omni_work_dir, "lib/"),
            "PATH": os.environ["PATH"] + ":" + os.path.join(self.omni_work_dir, "bin/"),
        }
        self.chatter_proc = subprocess.Popen(self.chat_server_cmd, env=env, cwd=self.omni_work_dir)

    @staticmethod
    def select_and_bind(socket: zmq.Socket) -> str:
        # randomly select a port between 1024 and 6553
        retry_count = 20
        err = None
        while retry_count > 0:
            try:
                port = random.randint(1024, 65535)
                # port = 5555
                url = f"tcp://localhost:{port}"
                socket.bind(url)
                return url
            except zmq.error.ZMQError as e:
                retry_count -= 1
                err = e
        raise err

    # immediate switch to status, discard prev_bytes, set immediate_switch to 1
    def immediate_switch_to(self, status):
        logger.warning(f"VA reader immediate switch to {status}")
        self.reset_prev = True
        self.status = status
        self.immediate_switch = 1
        if self.model_runner is not None:
            self.model_runner.pause_signal = True
            logger.warning(f"Model runner pause signal set to True")

    def recv_loop(self):
        while True:
            try:
                message = self.w2f_socket.recv()
            except Exception:
                logger.error(f"Error receiving message: {traceback.format_exc()}")
                break
            try:
                message = BSON.decode(message)
                msg_type = message["type"]
                logger.debug("Received message type: {}".format(msg_type))
                if msg_type == "AgentAudio":
                    audio = message["audio"]
                    if audio["type"] != "Pcm":
                        logger.error("Unsupported audio type: {}".format(audio["type"]))
                        continue
                    pcm_data = audio["data"]
                    audio_info = AudioInfo(audio["info"])
                    logger.debug("Received audio with duration: {}".format(audio_info.duration()))
                    if self.audio_info is None:
                        self.audio_info = audio_info
                    else:
                        # check if the audio info is the same
                        if not self.audio_info.is_spec_equal(audio_info):
                            raise ValueError("Audio info mismatch")
                    self.audio_buffer.add(pcm_data)
                    # if status is blank and has voice, set immediate switch to 1
                    if self.status == "blank" and self.has_voice(self.seg_duration):
                        self.immediate_switch_to("voice")
                elif msg_type == "AgentStartPlay":
                    logger.debug("Received AgentStartPlay, create new audio buffer")
                    self.audio_buffer = ByteBuffer()
                elif msg_type == "AgentEndPlay":
                    logger.debug("Received AgentEndPlay, mark audio finished")
                    self.audio_buffer.mark_finished()
                elif msg_type == "ClearAgentAudio":
                    logger.warning("Received ClearAgentAudio, clear audio buffer")
                    self.audio_buffer = None
                    self.audio_info = None
                    if self.status == "voice":
                        self.status = "blank"
                        # self.immediate_switch_to("blank")
            except Exception as e:
                logger.error("Error decoding message: {}, continue".format(e))
                continue
        logger.warning("recv loop interrupted")

    def start(self):
        self.launch_chat_server()
        self.recv_thread = threading.Thread(target=self.recv_loop)
        self.recv_thread.start()

    def has_voice(self, duration) -> bool:
        if self.audio_info is None or self.audio_buffer.current_size == 0:
            return False
        bytes_count = round(duration * self.audio_info.sample_rate) * self.audio_info.channel_count * 2  # S16LE assumed
        # if not has enough bytes and maybe has more voice, return False
        if self.audio_buffer.current_size < bytes_count and self.audio_buffer.has_more_voice():
            logger.warning(f"Not enough bytes and maybe has more voice, content_size: {self.audio_buffer.current_size}, bytes_count: {bytes_count}")
            return False
        return bytes_count

    def get_audio(self, fetch_duration) -> (bytes, AudioInfo):
        bytes_count = self.has_voice(fetch_duration)
        if bytes_count is False:
            return None
        pcm_data = self.audio_buffer.get(bytes_count)

        # the actual sample count fetched
        sample_count = len(pcm_data) // (self.audio_info.channel_count * 2)
        logger.debug("Fetched {} bytes audio".format(sample_count))
        logger.debug("After fetch, there are {} bytes left".format(self.audio_buffer.current_size))
        audio_info = deepcopy(self.audio_info)
        audio_info.sample_count = sample_count
        return (pcm_data, audio_info)

    def stop(self):
        self.model_runner = None
        if self.chatter_proc is not None:
            self.chatter_proc.terminate()
            self.chatter_proc.wait()
            self.chatter_proc = None
        self.w2f_socket.close()
        self.f2w_socket.close()

    def __del__(self):
        self.stop()


class OmniVAReader:
    def __init__(
        self,
        rank: int,
        world_size: int,
        stream_url: str,
        segment_duration: float = 5.0625,
        sample_rate: int = 16000,
        audio_channels: int = 1,
        buffer_size: int = 1,
        prev_duration: float = 0.3125,
        target_rank: int = 0,
        model_runner=None,
        huoshan_tts_voice_type=None,
    ):
        self.rank = rank
        self.world_size = world_size
        self.stream_url = stream_url
        self.segment_duration = segment_duration
        self.sample_rate = sample_rate

        self.audio_channels = audio_channels
        self.prev_duration = prev_duration
        self.all_seg_sample_count = int(self.segment_duration * self.sample_rate)
        self.prev_seg_sample_count = int(self.prev_duration * self.sample_rate)
        self.prev_seg_chunk = None

        self.target_rank = target_rank % self.world_size
        self.flag_tensor = torch.tensor([0], dtype=torch.int32).to(device="cuda")
        self.immediate_switch_tensor = torch.tensor([0], dtype=torch.int32).to(device="cuda")
        chunk_size = int(self.segment_duration * self.sample_rate) * 2
        self.audio_tensor = torch.zeros(chunk_size, dtype=torch.uint8, device="cuda")
        self.chat_adapter = None
        self.model_runner = model_runner
        self.huoshan_tts_voice_type = huoshan_tts_voice_type

        assert self.audio_channels == 1, "Only mono audio is supported for OmniVAReader"
        logger.info(f"VAReader initialized for stream: {stream_url} target_rank: {self.target_rank}")
        logger.info(f"Audio duration per chunk: {segment_duration}s, sample rate: {sample_rate}Hz")

    def init_omni_env(self):
        self.omni_work_dir = os.getenv("OMNI_WORK_DIR", "/path/of/seko_chatter/")
        self.session_id = os.getenv("OMNI_SESSION_ID", "")
        self.account = os.getenv("OMNI_ACCOUNT", "")
        self.config_files = os.getenv("OMNI_CONFIG_FILES", "").split(",")
        self.config_schema_path = os.getenv("OMNI_CONFIG_SCHEMA_PATH", None)
        assert os.path.exists(self.omni_work_dir), f"OMNI work directory {self.omni_work_dir} does not exist"
        assert self.session_id and self.account, "OMNI_SESSION_ID and OMNI_ACCOUNT are required"
        logger.info(
            f"OMNI work directory: {self.omni_work_dir}, session_id: {self.session_id}, account: {self.account}, config_files: {self.config_files}, config_schema_path: {self.config_schema_path}"
        )

    def start(self):
        if self.rank == self.target_rank:
            self.init_omni_env()
            assert self.stream_url.startswith("http"), "Only HTTP stream is supported for OmniVAReader"
            self.chat_adapter = ChatAdapter(
                omni_work_dir=self.omni_work_dir,
                whep_url=self.stream_url,
                session_id=self.session_id,
                account=self.account,
                config_files=self.config_files,
                config_schema_path=self.config_schema_path,
                seg_duration=self.segment_duration,
                model_runner=self.model_runner,
                huoshan_tts_voice_type=self.huoshan_tts_voice_type,
            )
            self.chat_adapter.start()
            logger.info(f"OmniVAReader {self.rank}/{self.world_size} started successfully")
        else:
            logger.info(f"OmniVAReader {self.rank}/{self.world_size} wait only")
        if self.world_size > 1:
            logger.info(f"OmniVAReader {self.rank}/{self.world_size} wait barrier")
            dist.barrier()
            logger.info(f"OmniVAReader {self.rank}/{self.world_size} end barrier")

    def braodcast_audio_data(self, audio_data):
        if self.rank == self.target_rank:
            if audio_data is None:
                self.flag_tensor.fill_(0)
            else:
                self.flag_tensor.fill_(1)
                self.audio_tensor.copy_(torch.frombuffer(bytearray(audio_data), dtype=torch.uint8))
                # logger.info(f"rank {self.rank} send audio_tensor: {self.audio_tensor.shape}")

        dist.broadcast(self.flag_tensor, src=self.target_rank)
        if self.flag_tensor.item() == 0:
            return None

        dist.broadcast(self.audio_tensor, src=self.target_rank)
        if self.rank != self.target_rank:
            # logger.info(f"rank {self.rank} recv audio_tensor: {self.audio_tensor.shape}")
            audio_data = self.audio_tensor.cpu().numpy().tobytes()
        return audio_data

    def bytes_to_ndarray(self, audio_data):
        if audio_data is None:
            return None
        audio_data = np.frombuffer(audio_data, dtype=np.int16)
        audio_data = audio_data.astype(np.float32) / 32768.0
        # logger.info(f"Got segment audio rank={self.rank}: {audio_data.shape} {audio_data.dtype} {audio_data.min()} {audio_data.max()}")
        return audio_data

    def convert_pcm_s16le_to_mono_resampled(self, audio_data, audio_info):
        audio = np.frombuffer(audio_data, dtype=np.int16)
        sample_count = audio_info.sample_count
        assert len(audio) == sample_count * audio_info.channel_count, f"audio length {len(audio)} != sample_count * channel_count {sample_count * audio_info.channel_count}"
        # convert to mono
        if audio_info.channel_count > 1:
            audio = audio.reshape(-1, audio_info.channel_count).mean(axis=1)

        # logger.info(f"audio: {audio.shape} {audio.dtype} {audio.min()} {audio.max()}")
        if audio_info.sample_rate != self.sample_rate:
            sample_count = int(len(audio) * self.sample_rate / audio_info.sample_rate)
            audio = resample(audio, sample_count).astype(np.int16)
            # logger.info(f"resampled audio: {audio.shape} {audio.dtype} {audio.min()} {audio.max()} {sample_count}")
        logger.warning(f"valid audio: {audio.shape} {audio.dtype} {audio.min()} {audio.max()} {sample_count}")
        return audio, sample_count

    def prepare_audio_data(self, chat_audio_result):
        sample_count = 0
        audio = np.array([], dtype=np.int16)

        # convert chat audio result to mono and target sample rate
        if chat_audio_result is not None:
            audio_data, audio_info = chat_audio_result
            audio, sample_count = self.convert_pcm_s16le_to_mono_resampled(audio_data, audio_info)

        # if is not the first segment, concat with previous segment
        if self.prev_seg_chunk is not None:
            audio = np.concatenate([self.prev_seg_chunk, audio])
            sample_count = len(audio)
        assert sample_count <= self.all_seg_sample_count, f"audio length {sample_count} > all_seg_sample_count {self.all_seg_sample_count}"

        # pad 0 to the audio to make it the same length as all_seg_sample_count
        if sample_count < self.all_seg_sample_count:
            pad_count = self.all_seg_sample_count - sample_count
            # logger.info(f"pad {pad_count} samples to audio")
            audio = np.pad(audio, (0, pad_count), mode="constant", constant_values=0)
            sample_count = len(audio)

        # update prev seg chunk
        self.prev_seg_chunk = audio[-self.prev_seg_sample_count :]
        # logger.info(f"audio: {audio.shape} {audio.dtype} {audio.min()} {audio.max()} {sample_count}, prev seg chunk: {self.prev_seg_chunk.shape}")
        return audio.tobytes()

    def get_fetch_duration(self):
        fetch_duration = self.segment_duration
        # after immediate switch, reset prev seg chunk
        if self.chat_adapter.reset_prev:
            self.prev_seg_chunk = None
            self.chat_adapter.reset_prev = False
            logger.warning(f"Reset prev seg chunk")
        # first segment, fetch segment_duration, else fetch segment_duration - prev_duration
        if self.prev_seg_chunk is not None:
            fetch_duration -= self.prev_duration
        return fetch_duration

    def get_audio_segment(self):
        audio_data = None
        if self.rank == self.target_rank:
            try:
                fetch_duration = self.get_fetch_duration()
                # logger.info(f"Get segment, fetch_duration: {fetch_duration}")
                if self.chat_adapter.status == "voice":
                    audio_result = self.chat_adapter.get_audio(fetch_duration)
                    audio_data = self.prepare_audio_data(audio_result)
                    # think all voice segments inferred, naturally switch to blank
                    if audio_result is None:
                        logger.info(f"Think all voice segments inferred, naturally switch to blank")
                        self.chat_adapter.status = "blank"
                else:
                    audio_data = self.prepare_audio_data(None)
            except Exception as e:
                logger.warning(f"Failed to get voice segment: {e}")
                return None
        if self.world_size > 1:
            audio_data = self.braodcast_audio_data(audio_data)
        audio_data = self.bytes_to_ndarray(audio_data)
        return audio_data

    def get_immediate_switch(self):
        if self.rank == self.target_rank:
            if self.chat_adapter.immediate_switch == 1:
                self.immediate_switch_tensor.fill_(1)
                # reset immediate switch
                self.chat_adapter.immediate_switch = 0
            else:
                self.immediate_switch_tensor.fill_(0)
        dist.broadcast(self.immediate_switch_tensor, src=self.target_rank)
        immediate_switch = self.immediate_switch_tensor.item()
        return immediate_switch

    def stop(self):
        self.model_runner = None
        if self.chat_adapter is not None:
            self.chat_adapter.stop()
            self.chat_adapter = None
            logger.warning("OmniVAReader stopped")

    def __del__(self):
        self.stop()


if __name__ == "__main__":
    WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
    RANK = int(os.environ.get("RANK", 0))
    if WORLD_SIZE > 1:
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(dist.get_rank())
        logger.info(f"Distributed initialized: rank={RANK}, world_size={WORLD_SIZE}")

    reader = OmniVAReader(
        RANK,
        WORLD_SIZE,
        "https://reverse.st-oc-01.chielo.org/10.5.64.49:8000/rtc/v1/whep/?app=publish&stream=test_stream_ll&eip=10.120.114.82:8000",
        segment_duration=17 / 16,
        sample_rate=16000,
        audio_channels=1,
        prev_duration=1 / 16,
    )
    reader.start()
    fail_count = 0
    max_fail_count = 100000000

    try:
        while True:
            audio_data = reader.get_audio_segment(timeout=1)
            if audio_data is not None:
                logger.info(f"Got audio chunk, shape: {audio_data.shape}, range: [{audio_data.min()}, {audio_data.max()}]")
                fail_count = 0
            else:
                fail_count += 1
                if fail_count > max_fail_count:
                    logger.warning("Failed to get audio chunk, stop reader")
                    reader.stop()
                    break
            time.sleep(0.95)
    finally:
        reader.stop()