数字人实时对话AvatarDialog API参考

import asyncio
import hashlib
import os
import websockets
import base64
import json
import wave
import time
import uuid
import argparse
from copy import deepcopy
from loguru import logger
import numpy as np
import struct


parser = argparse.ArgumentParser()
parser.add_argument("--loop_num", type=int, default=1, help="loop_num")
parser.add_argument("--concurrent_num", type=int, default=1, help="concurrent_num")
parser.add_argument("--log_file_name", type=str, default="default.log", help="log_file_name")
args = parser.parse_args()

DS_WSS_ENDPOINT = "wss://dashscope.aliyuncs.com/api-ws/v1/inference"
MODEL = "avatar-dialog"
AVATAR_ID = "jiaoyue"
DS_AUTH = "Bearer " + os.getenv("DASHSCOPE_API_KEY")

AUDIO_PATH = "wav/audio48k.wav"	# 测试音频路径
TEST_TIME = 3600 * 8  # 测试时间 单位秒

RTC_H5_URL = "https://market.wapa.taobao.com/app/xr-paas/realtime-avatar-preview/index.html#/rtc"

AUDIO_COPY_TIMES = 1
AUDIO_SLICE_RATE = 2
sample_rata = 16000

# MAX_AUDIO_DURATION = 10
MAX_AUDIO_DURATION = -1

LOG_FILE_NAME = args.log_file_name  # 日志名

logger.add(LOG_FILE_NAME, rotation="1 week")

CHANNEL_ID = "test_channel"


def get_rtc_param(user_id):
    app_id = os.getenv("RTC_APP_ID")
    app_key = os.getenv("RTC_APP_KEY")
    timestamp = int(time.time()) + 24 * 60 * 60
    sha256_hash = hashlib.sha256()
    data = f"{app_id}{app_key}{CHANNEL_ID}{user_id}{timestamp}"
    sha256_hash.update(data.encode())
    token = sha256_hash.hexdigest()
    return {
        "app_id": app_id,
        "channel_id": CHANNEL_ID,
        "user_id": user_id,
        "nonce": "",
        "timestamp": timestamp,
        "token": token,
        "gslb": ["https://rgslb.rtc.aliyuncs.com"]
    }


def get_subscribe_url():
    rtc_param = get_rtc_param(user_id=str(uuid.uuid4()))
    url = (f"{RTC_H5_URL}?userId={rtc_param['user_id']}"
           f"&channelId={CHANNEL_ID}"
           f"&appId={rtc_param['app_id']}"
           f"&token={rtc_param['token']}"
           f"&timestamp={rtc_param['timestamp']}"
           f"&userName={str(uuid.uuid4())}"
          )
    return url


def get_init_msg(task_id):
    msg_initialize_video_session = {
        "header": {
            "task_id": task_id,
            "streaming": "duplex",
            "action": "run-task"
        },
        "payload": {
            "task_group": "aigc",
            "task": "video-generation",
            "function": "stream-generation",
            "model": MODEL,
            "input": {
                "header": {
                    "name": "InitializeVideoSession",
                    "request_id": str(uuid.uuid4())
                },
                "payload": {
                    "avatar_id": AVATAR_ID,
                    "format": "PCM",
                    "sample_rate": sample_rata,
                    "session_config": {
                        "enable_motion_data": False,
                        "motion_data_compression": "zip",
                        "motion_data_audio_encoder": "opus"
                    },
                    "h5_mode_enable": True,
                    "user_agent": {
                        "client": "BAILIAN",
                        "platform": "Dalvik/2.1.0 (Linux; U; Android 12; NOH-AL00 Build/HUAWEINOH-AL00)",
                        "version": "3.18.0",
                        "app_type": "Dev"
                    },
                    "rtc_param": get_rtc_param("test_user")
                }
            },
            "parameter": {}
        }
    }
    return json.dumps(msg_initialize_video_session)


def get_alive_message(task_id):
    msg_destroy_video_session = {
        "header": {
            "task_id": task_id,
            "streaming": "duplex",
            "action": "continue-task"
        },
        "payload": {
            "input": {
                "header": {
                    "name": "TriggerHeartbeat",
                    "request_id": str(uuid.uuid4())
                },
                "payload": {}
            }
        }
    }
    return json.dumps(msg_destroy_video_session)


def get_destroy_message(task_id):
    msg_destroy_video_session = {
        "header": {
            "task_id": task_id,
            "streaming": "duplex",
            "action": "finish-task"
        },
        "payload": {
            "input": {
                "header": {
                    "name": "DestroyVideoSession",
                    "request_id": str(uuid.uuid4())
                },
                "payload": {}
            }
        }
    }
    return json.dumps(msg_destroy_video_session)


def get_send_speech_text_message(task_id):
    send_speech_text_message = {
        "header": {
            "task_id": task_id,
            "streaming": "duplex",
            "action": "continue-task"
        },
        "payload": {
            "input": {
                "header": {
                    "name": "SendSpeechText",
                    "request_id": str(uuid.uuid4())
                },
                "payload": {
                    "task_id": task_id,
                    "source": "llm",
                    "mode": "full_text",
                    "text": "你好，我是云小宝",
                    "end_of_task": True
                }
            }
        }
    }
    return json.dumps(send_speech_text_message)


def get_inner_msg(msg):
    msg_json = json.loads(msg)
    output = msg_json["payload"].get("output", None)
    usage_msg = msg_json["payload"].get("usage", None)
    return output, usage_msg


def read_wav(wav_path: str):
    with wave.open(wav_path, 'rb') as wav_file:
        frame_rate = wav_file.getframerate()
        num_frames = wav_file.getnframes()
        audio_data = wav_file.readframes(num_frames)
        if MAX_AUDIO_DURATION > 0:
            audio_data = audio_data[:int(frame_rate * MAX_AUDIO_DURATION * 2)]
        return audio_data, frame_rate

def read_pcm_to_bytes(file_path: str, sample_rate):
    with open(file_path, 'rb') as f:
        audio_data = f.read()
    return audio_data, sample_rate

def read_pcm_file(file_path, sample_rate, sample_width=2, num_channels=1):
    """
    读取 PCM 格式的音频文件。

    :param file_path: PCM 文件的路径。
    :param sample_rate: PCM 音频的采样率。
    :param sample_width: 每个音频样本的字节数（例如，16 位音频对应2字节）。
    :param num_channels: 音频的通道数（例如，单声道是1，立体声是2）。
    :return: 一个包含音频数据的 numpy 数组，以及采样率。
    """
    # 打开 PCM 文件读取二进制数据
    with open(file_path, 'rb') as pcm_file:
        # 读取所有数据
        pcm_data = pcm_file.read()

    # 确定每个样本的格式
    if sample_width == 1:
        sample_format = 'b'  # 8-bit PCM
    elif sample_width == 2:
        sample_format = 'h'  # 16-bit PCM
    else:
        raise ValueError("Unsupported sample width: {}".format(sample_width))

    # 计算样本数量
    num_samples = len(pcm_data) // (sample_width * num_channels)

    # 使用 struct.unpack 解码二进制数据到整数
    audio_data = struct.unpack('<' + sample_format * num_samples * num_channels, pcm_data)

    # 将数据转换为 numpy 数组
    audio_array = np.array(audio_data, dtype=np.int16)

    # 如果是多通道音频，将数据重塑为（num_samples, num_channels）的形状
    if num_channels > 1:
        audio_array = audio_array.reshape(-1, num_channels)

    return audio_array, sample_rate


def get_audio_msg_list(task_id):
    audio_path = AUDIO_PATH
    audio_file, audio_sample_rate = read_wav(audio_path)
    # audio_file, audio_sample_rate = read_pcm_to_bytes(audio_path, 16000)
    global sample_rata
    sample_rata = audio_sample_rate
    audio_file = audio_file * AUDIO_COPY_TIMES
    end_of_task = False
    start = 0
    msg_list = []

    while not end_of_task:
        end = int(start + audio_sample_rate * AUDIO_SLICE_RATE)
        audio_data = audio_file[start:end]
        end_of_task = end >= len(audio_file) - 1
        audio_base64 = base64.b64encode(audio_data).decode("utf8")
        request_id = str(uuid.uuid4())
        msg_generate_video = {
            "header": {
                "task_id": task_id,
                "streaming": "duplex",
                "action": "continue-task"
            },
            "payload": {
                "input": {
                    "header": {
                        "name": "GenerateVideo",
                        "request_id": request_id
                    },
                    "payload": {
                        "speech_id": "",
                        "audio_data": "",
                        "end_of_task": end_of_task
                    }
                }
            }
        }
        msg_generate_video["payload"]["input"]["payload"]["audio_data"] = audio_base64
        msg_generate_video["payload"]["input"]["payload"]["end_of_task"] = end_of_task
        msg_generate_video["payload"]["input"]["payload"]["end_of_speech"] = end_of_task
        msg_list.append(msg_generate_video)
        if end_of_task:
            break
        else:
            start = end
    return msg_list


def update_msg_speech_id(msg_list):
    task_id = str(uuid.uuid4())
    for msg in msg_list:
        msg["payload"]["input"]["payload"]["task_id"] = task_id
        msg["payload"]["input"]["payload"]["speech_id"] = task_id


def update_msg_request_id(msg_list):
    request_id = str(uuid.uuid4())
    for msg in msg_list:
        msg["payload"]["input"]["header"]["request_id"] = request_id


class TimeCost:
    def __init__(self):
        self.init_start = None
        self.first_audio_data_send_ts = None
        self.first_motion_data_recv_ts = None


async def init(worker_num, ws, rec_timer):
    while True:
        msg = await ws.recv()
        logger.info("{}", msg)
        inner_msg, usage_msg = get_inner_msg(msg)
        if inner_msg and inner_msg["header"]["name"] == "VideoSessionInitialized":
            latency = int((time.time() - rec_timer.init_start) * 1000)
            logger.info(f"worker {worker_num} =================== initCost(in ms): {latency}")
            try:
                rtc_param = inner_msg["payload"]["rtc_param"]
                channel_id = rtc_param["channel_id"]
                app_id = rtc_param["app_id"]
                user_id = rtc_param["user_id"]
                token = rtc_param["token"]
                timestamp = rtc_param["timestamp"]

                assert app_id != ""

                url = (f"{RTC_H5_URL}?userId={user_id}"
                    f"&channelId={channel_id}"
                    f"&appId={app_id}"
                    f"&token={token}"
                    f"&timestamp={timestamp}"
                    f"&userName={str(uuid.uuid4())}"
                    )
            except:
                url = get_subscribe_url()
            print(f"Opening rtc channel {url}")
        elif inner_msg and inner_msg["header"]["name"] == "VideoSessionStarted":
            latency = int((time.time() - rec_timer.init_start) * 1000)
            logger.info(f"worker {worker_num} =================== startCost(in ms): {latency}")
            return


async def destroy(worker_num, ws, rec_timer):
    while True:
        msg = await ws.recv()
        inner_msg, usage_msg = get_inner_msg(msg)
        logger.info("receive message: {}", inner_msg)
        if inner_msg and inner_msg["header"]["name"] == "VideoSessionDestroyed":
            logger.info(f"worker {worker_num} =================== VideoSessionDestroyed")
            logger.info("任务结束，查看计量: {}", usage_msg)
            return


async def send_audio_messages(worker_num, ws, msg_list, rec_timer, alive_msg):
    for j in range(int(TEST_TIME / 10)):
        uu_msg = deepcopy(msg_list)
        update_msg_speech_id(uu_msg)
        update_msg_request_id(uu_msg)
        total_audio_duration = 0
        start_send_time = time.time()
        for i, message in enumerate(uu_msg):
            await ws.send(json.dumps(message))
            if i == 0:
                speech_id = message["payload"]["input"]["payload"]["task_id"]
                logger.info(
                    f"worker {worker_num} first send_audio_messages cost: {int((time.time() - rec_timer.first_audio_data_send_ts) * 1000)} speech_id: {speech_id}")
                rec_timer.first_audio_data_send_ts = time.time()
            current_audio_duration = AUDIO_SLICE_RATE / 2
            total_audio_duration += current_audio_duration
            sleep_time = current_audio_duration / 2
            await asyncio.sleep(sleep_time)
        time_to_sleep = (total_audio_duration - (time.time() - start_send_time)) + 5
        logger.info("total audio duration: {}, sleep time: {}", total_audio_duration, time_to_sleep)
        await asyncio.sleep(time_to_sleep)
        await ws.send(alive_msg)


async def demo(uri, auth_token, audio_msg_list, timer):
    global ws
    task_id = str(uuid.uuid4())
    try:
        timer = TimeCost()
        worker_num = 0
        async with websockets.connect(uri, ping_interval=300, ping_timeout=300,
                                      additional_headers={"Authorization": auth_token}, open_timeout=300) as websocket:
            ws = websocket
            init_message = get_init_msg(task_id)
            await websocket.send(init_message)
            timer.init_start = time.time()
            logger.info(f"demo init_start")
            await init(worker_num, websocket, timer)

            send_speech_text_message = get_send_speech_text_message(task_id)
            await websocket.send(send_speech_text_message)

            timer.first_audio_data_send_ts = time.time()
            alive_msg = get_alive_message(task_id)
            send_task = asyncio.create_task(
                send_audio_messages(worker_num, websocket, audio_msg_list, timer, alive_msg))

            # create destroy loop to recv message
            destroy_task = asyncio.create_task(
                destroy(worker_num, websocket, timer)
            )

            await send_task

            destroy_message = get_destroy_message(task_id)
            await websocket.send(destroy_message)
            logger.info(f"worker {worker_num}, sent destroy msg: {destroy_message}")
            
            await destroy_task

    except Exception as e:
        logger.opt(exception=True).error("{}", e)
    finally:
        if ws:
            await close_ws(ws, task_id)


async def async_only():
    audio_msg_list = get_audio_msg_list(str(uuid.uuid4()))
    timer = TimeCost()
    await demo(DS_WSS_ENDPOINT, DS_AUTH, audio_msg_list, timer)


async def close_ws(websocket, task_id):
    try:
        logger.info("ws state: {}", websocket.state)
        if websocket.state != 1:
            return
        destroy_message = get_destroy_message(task_id)
        await websocket.send(destroy_message)
        logger.info(f"worker {0}, sent destroy msg: {destroy_message}")
        await destroy(0, websocket, None)
    except Exception as e:
        logger.opt(exception=True).error("Failed to close WebSocket: {}", e)


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(async_only())
    except KeyboardInterrupt:
        print("KeyboardInterrupt, stop")
        if ws:
            loop.run_until_complete(close_ws(ws))
        # 取消所有任务
        tasks = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task(loop)]
        for task in tasks:
            task.cancel()
        loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
        print("All tasks cancelled.")
    finally:
        loop.stop()
        loop.close()