import base64
import ctypes
import io
import json
import os
import struct
from dataclasses import dataclass
from enum import Enum
from typing import AsyncGenerator, Union
import httpx
import numpy as np
import ormsgpack
import soundfile as sf
from .schema import (
ServeMessage,
ServeRequest,
ServeTextPart,
ServeVQGANDecodeRequest,
ServeVQGANEncodeRequest,
ServeVQPart,
)
class CustomAudioFrame:
def __init__(self, data, sample_rate, num_channels, samples_per_channel):
if len(data) < num_channels * samples_per_channel * ctypes.sizeof(
ctypes.c_int16
):
raise ValueError(
"data length must be >= num_channels * samples_per_channel * sizeof(int16)"
)
self._data = bytearray(data)
self._sample_rate = sample_rate
self._num_channels = num_channels
self._samples_per_channel = samples_per_channel
@property
def data(self):
return memoryview(self._data).cast("h")
@property
def sample_rate(self):
return self._sample_rate
@property
def num_channels(self):
return self._num_channels
@property
def samples_per_channel(self):
return self._samples_per_channel
@property
def duration(self):
return self.samples_per_channel / self.sample_rate
def __repr__(self):
return (
f"CustomAudioFrame(sample_rate={self.sample_rate}, "
f"num_channels={self.num_channels}, "
f"samples_per_channel={self.samples_per_channel}, "
f"duration={self.duration:.3f})"
)
class FishE2EEventType(Enum):
SPEECH_SEGMENT = 1
TEXT_SEGMENT = 2
END_OF_TEXT = 3
END_OF_SPEECH = 4
ASR_RESULT = 5
USER_CODES = 6
@dataclass
class FishE2EEvent:
type: FishE2EEventType
frame: np.ndarray = None
text: str = None
vq_codes: list[list[int]] = None
client = httpx.AsyncClient(
timeout=None,
limits=httpx.Limits(
max_connections=None,
max_keepalive_connections=None,
keepalive_expiry=None,
),
)
class FishE2EAgent:
def __init__(self):
self.llm_url = "http://localhost:8080/v1/chat"
self.vqgan_url = "http://localhost:8080"
self.client = httpx.AsyncClient(timeout=None)
async def get_codes(self, audio_data, sample_rate):
audio_buffer = io.BytesIO()
sf.write(audio_buffer, audio_data, sample_rate, format="WAV")
audio_buffer.seek(0)
# Step 1: Encode audio using VQGAN
encode_request = ServeVQGANEncodeRequest(audios=[audio_buffer.read()])
encode_request_bytes = ormsgpack.packb(
encode_request, option=ormsgpack.OPT_SERIALIZE_PYDANTIC
)
encode_response = await self.client.post(
f"{self.vqgan_url}/v1/vqgan/encode",
data=encode_request_bytes,
headers={"Content-Type": "application/msgpack"},
)
encode_response_data = ormsgpack.unpackb(encode_response.content)
codes = encode_response_data["tokens"][0]
return codes
async def stream(
self,
system_audio_data: np.ndarray | None,
user_audio_data: np.ndarray | None,
sample_rate: int,
num_channels: int,
chat_ctx: dict | None = None,
) -> AsyncGenerator[bytes, None]:
if system_audio_data is not None:
sys_codes = await self.get_codes(system_audio_data, sample_rate)
else:
sys_codes = None
if user_audio_data is not None:
user_codes = await self.get_codes(user_audio_data, sample_rate)
# Step 2: Prepare LLM request
if chat_ctx is None:
sys_parts = [
ServeTextPart(
text='您是由 Fish Audio 设计的语音助手，提供端到端的语音交互，实现无缝用户体验。首先转录用户的语音，然后使用以下格式回答："Question: [用户语音]\n\nAnswer: [你的回答]\n"。'
),
]
if system_audio_data is not None:
sys_parts.append(ServeVQPart(codes=sys_codes))
chat_ctx = {
"messages": [
ServeMessage(
role="system",
parts=sys_parts,
),
],
}
else:
if chat_ctx["added_sysaudio"] is False and sys_codes:
chat_ctx["added_sysaudio"] = True
chat_ctx["messages"][0].parts.append(ServeVQPart(codes=sys_codes))
prev_messages = chat_ctx["messages"].copy()
if user_audio_data is not None:
yield FishE2EEvent(
type=FishE2EEventType.USER_CODES,
vq_codes=user_codes,
)
else:
user_codes = None
request = ServeRequest(
messages=prev_messages
+ (
[
ServeMessage(
role="user",
parts=[ServeVQPart(codes=user_codes)],
)
]
if user_codes
else []
),
streaming=True,
num_samples=1,
)
# Step 3: Stream LLM response and decode audio
buffer = b""
vq_codes = []
current_vq = False
async def decode_send():
nonlocal current_vq
nonlocal vq_codes
data = np.concatenate(vq_codes, axis=1).tolist()
# Decode VQ codes to audio
decode_request = ServeVQGANDecodeRequest(tokens=[data])
decode_response = await self.client.post(
f"{self.vqgan_url}/v1/vqgan/decode",
data=ormsgpack.packb(
decode_request,
option=ormsgpack.OPT_SERIALIZE_PYDANTIC,
),
headers={"Content-Type": "application/msgpack"},
)
decode_data = ormsgpack.unpackb(decode_response.content)
# Convert float16 audio data to int16
audio_data = np.frombuffer(decode_data["audios"][0], dtype=np.float16)
audio_data = (audio_data * 32768).astype(np.int16).tobytes()
audio_frame = CustomAudioFrame(
data=audio_data,
samples_per_channel=len(audio_data) // 2,
sample_rate=44100,
num_channels=1,
)
yield FishE2EEvent(
type=FishE2EEventType.SPEECH_SEGMENT,
frame=audio_frame,
vq_codes=data,
)
current_vq = False
vq_codes = []
async with self.client.stream(
"POST",
self.llm_url,
data=ormsgpack.packb(request, option=ormsgpack.OPT_SERIALIZE_PYDANTIC),
headers={"Content-Type": "application/msgpack"},
) as response:
async for chunk in response.aiter_bytes():
buffer += chunk
while len(buffer) >= 4:
read_length = struct.unpack("I", buffer[:4])[0]
if len(buffer) < 4 + read_length:
break
body = buffer[4 : 4 + read_length]
buffer = buffer[4 + read_length :]
data = ormsgpack.unpackb(body)
if data["delta"] and data["delta"]["part"]:
if current_vq and data["delta"]["part"]["type"] == "text":
async for event in decode_send():
yield event
if data["delta"]["part"]["type"] == "text":
yield FishE2EEvent(
type=FishE2EEventType.TEXT_SEGMENT,
text=data["delta"]["part"]["text"],
)
elif data["delta"]["part"]["type"] == "vq":
vq_codes.append(np.array(data["delta"]["part"]["codes"]))
current_vq = True
if current_vq and vq_codes:
async for event in decode_send():
yield event
yield FishE2EEvent(type=FishE2EEventType.END_OF_TEXT)
yield FishE2EEvent(type=FishE2EEventType.END_OF_SPEECH)
# Example usage:
async def main():
import torchaudio
agent = FishE2EAgent()
# Replace this with actual audio data loading
with open("uz_story_en.m4a", "rb") as f:
audio_data = f.read()
audio_data, sample_rate = torchaudio.load("uz_story_en.m4a")
audio_data = (audio_data.numpy() * 32768).astype(np.int16)
stream = agent.stream(audio_data, sample_rate, 1)
if os.path.exists("audio_segment.wav"):
os.remove("audio_segment.wav")
async for event in stream:
if event.type == FishE2EEventType.SPEECH_SEGMENT:
# Handle speech segment (e.g., play audio or save to file)
with open("audio_segment.wav", "ab+") as f:
f.write(event.frame.data)
elif event.type == FishE2EEventType.ASR_RESULT:
print(event.text, flush=True)
elif event.type == FishE2EEventType.TEXT_SEGMENT:
print(event.text, flush=True, end="")
elif event.type == FishE2EEventType.END_OF_TEXT:
print("\nEnd of text reached.")
elif event.type == FishE2EEventType.END_OF_SPEECH:
print("End of speech reached.")
if __name__ == "__main__":
import asyncio
asyncio.run(main())