WebSocket 驱动数字人说话

Go

package main

import (
	"bytes"
	"crypto/md5"
	"crypto/tls"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

// 配置信息
const (
	AppID        = 123456               // 替换为您的AppID
	ServerSecret = "your_server_secret" // 替换为您的ServerSecret
	TaskID       = "your_task_id"       // 替换为您的数字人视频流任务ID
	APIHost      = "aigc-digitalhuman-api.zegotech.cn"
	PCMFilePath  = "./audio.pcm" // 替换为您的PCM音频文件路径
)

// DriveByWsStreamResponse 获取WebSocket驱动信息的响应
type DriveByWsStreamResponse struct {
	Code      int    `json:"Code"`
	Message   string `json:"Message"`
	RequestID string `json:"RequestId"`
	Data      struct {
		DriveID    string `json:"DriveId"`
		WssAddress string `json:"WssAddress"`
	} `json:"Data"`
}

// StartCommand WebSocket开始驱动指令
type StartCommand struct {
	Action  string `json:"Action"`
	Payload struct {
		DriveID    string `json:"DriveId"`
		SampleRate int    `json:"SampleRate"`
	} `json:"Payload"`
}

// StopCommand WebSocket停止驱动指令
type StopCommand struct {
	Action  string `json:"Action"`
	Payload struct {
		DriveID string `json:"DriveId"`
	} `json:"Payload"`
}

func generateSignature(appId uint32, signatureNonce string, serverSecret string, timestamp int64) (Signature string) {
	data := fmt.Sprintf("%d%s%s%d", appId, signatureNonce, serverSecret, timestamp)
	h := md5.New()
	h.Write([]byte(data))
	return hex.EncodeToString(h.Sum(nil))
}

// 获取WebSocket驱动信息
func getDriveByWsStreamInfo() (*DriveByWsStreamResponse, error) {
	// 构造请求参数
	signatureNonce := strconv.FormatInt(rand.Int63(), 10)
	timestamp := time.Now().Unix()
	params := map[string]string{
		"Action":           "DriveByWsStream",
		"AppId":            strconv.FormatInt(AppID, 10),
		"SignatureNonce":   signatureNonce,
		"SignatureVersion": "2.0",
		"Timestamp":        strconv.FormatInt(timestamp, 10),
	}

	// 生成签名
	signature := generateSignature(AppID, signatureNonce, ServerSecret, timestamp)
	params["Signature"] = signature

	// 构造请求URL
	baseURL := fmt.Sprintf("https://%s/", APIHost)
	reqURL, err := url.Parse(baseURL)
	if err != nil {
		return nil, err
	}

	// 添加查询参数
	query := reqURL.Query()
	for k, v := range params {
		query.Set(k, v)
	}
	reqURL.RawQuery = query.Encode()

	// 构造请求体
	reqBody := map[string]string{
		"TaskId": TaskID,
	}
	reqBodyBytes, err := json.Marshal(reqBody)
	if err != nil {
		return nil, err
	}

	// 发送HTTP请求
	req, err := http.NewRequest("POST", reqURL.String(), bytes.NewBuffer(reqBodyBytes))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// 解析响应
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var response DriveByWsStreamResponse
	err = json.Unmarshal(body, &response)
	if err != nil {
		return nil, err
	}

	if response.Code != 0 {
		return nil, fmt.Errorf("API error: %d - %s", response.Code, response.Message)
	}

	return &response, nil
}

// 使用WebSocket驱动数字人说话
func driveDigitalHumanWithWebSocket(wssAddress, driveID string, pcmFilePath string) error {
	// 1. 建立WebSocket连接
	dialer := websocket.DefaultDialer
	dialer.TLSClientConfig = &tls.Config{
		InsecureSkipVerify: true, // 忽略证书验证
	}
	c, _, err := dialer.Dial(wssAddress, nil)
	if err != nil {
		return fmt.Errorf("连接WebSocket失败: %v", err)
	}
	defer c.Close()

	// 2. 发送Start开始驱动指令
	startCmd := StartCommand{
		Action: "Start",
		Payload: struct {
			DriveID    string `json:"DriveId"`
			SampleRate int    `json:"SampleRate"`
		}{
			DriveID:    driveID,
			SampleRate: 16000, // 使用16000Hz采样率
		},
	}

	startCmdBytes, err := json.Marshal(startCmd)
	if err != nil {
		return fmt.Errorf("序列化Start指令失败: %v", err)
	}

	err = c.WriteMessage(websocket.TextMessage, startCmdBytes)
	if err != nil {
		return fmt.Errorf("发送Start指令失败: %v", err)
	}

	// 3. 读取PCM文件并发送音频数据
	pcmFile, err := os.Open(pcmFilePath)
	if err != nil {
		return fmt.Errorf("打开PCM文件失败: %v", err)
	}
	defer pcmFile.Close()

	// 每次读取并发送4KB的PCM数据
	buffer := make([]byte, 4096)
	for {
		n, err := pcmFile.Read(buffer)
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("读取PCM文件失败: %v", err)
		}

		// 发送PCM二进制数据
		err = c.WriteMessage(websocket.BinaryMessage, buffer[:n])
		if err != nil {
			return fmt.Errorf("发送PCM数据失败: %v", err)
		}

		// 控制发送速率，避免发送过快
		time.Sleep(100 * time.Millisecond)
	}

	// 4. 发送Stop停止驱动指令
	stopCmd := StopCommand{
		Action: "Stop",
		Payload: struct {
			DriveID string `json:"DriveId"`
		}{
			DriveID: driveID,
		},
	}

	stopCmdBytes, err := json.Marshal(stopCmd)
	if err != nil {
		return fmt.Errorf("序列化Stop指令失败: %v", err)
	}

	err = c.WriteMessage(websocket.TextMessage, stopCmdBytes)
	if err != nil {
		return fmt.Errorf("发送Stop指令失败: %v", err)
	}

	return nil
}

func main() {
	// 1. 获取WebSocket驱动信息
	fmt.Println("正在获取WebSocket驱动信息...")
	response, err := getDriveByWsStreamInfo()
	if err != nil {
		fmt.Printf("获取WebSocket驱动信息失败: %v\n", err)
		return
	}

	fmt.Printf("获取WebSocket驱动信息成功:\n")
	fmt.Printf("DriveId: %s\n", response.Data.DriveID)
	fmt.Printf("WssAddress: %s\n", response.Data.WssAddress)

	// 2. 使用WebSocket驱动数字人说话
	fmt.Println("正在使用WebSocket驱动数字人说话...")
	err = driveDigitalHumanWithWebSocket(response.Data.WssAddress, response.Data.DriveID, PCMFilePath)
	if err != nil {
		fmt.Printf("WebSocket驱动数字人说话失败: %v\n", err)
		return
	}

	fmt.Println("WebSocket驱动数字人说话成功!")
}