怎么指定说话人啊 ，每次都是随机一个说话人

[anan1213095357]

from tools.vqgan.inference import load_model as load_decoder_model
from tools.llama.generate import (
    GenerateRequest,
    GenerateResponse,
    WrappedGenerateResponse,
    launch_thread_safe_queue,
)
from tools.file import AUDIO_EXTENSIONS, list_files
from tools.api import decode_vq_tokens, encode_reference
from fish_speech.utils import autocast_exclude_mps, set_seed
from fish_speech.text.chn_text_norm.text import Text as ChnNormedText
from fish_speech.i18n import i18n
import gc
import html
import io
import os
import queue
import wave
from argparse import ArgumentParser
from functools import partial
from pathlib import Path

import librosa
import numpy as np
import pyrootutils
import torch
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import StreamingResponse, FileResponse, JSONResponse
from loguru import logger
from pydantic import BaseModel
from transformers import AutoTokenizer

pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)


# Make einx happy
os.environ["EINX_FILTER_TRACEBACK"] = "false"

app = FastAPI()


def build_html_error_message(error):
    return f"""
    <div style="color: red;
    font-weight: bold;">
        {html.escape(str(error))}
    </div>
    """


@torch.inference_mode()
def inference(
    text,
    enable_reference_audio,
    reference_audio,
    reference_text,
    max_new_tokens,
    chunk_length,
    top_p,
    repetition_penalty,
    temperature,
    seed="0",
    streaming=False,
):
    if args.max_gradio_length > 0 and len(text) > args.max_gradio_length:
        return (
            None,
            None,
            i18n("Text is too long, please keep it under {} characters.").format(
                args.max_gradio_length
            ),
        )

    seed = int(seed)
    if seed != 0:
        set_seed(seed)
        logger.warning(f"set seed: {seed}")

    # Parse reference audio aka prompt
    prompt_tokens = encode_reference(
        decoder_model=decoder_model,
        reference_audio=reference_audio,
        enable_reference_audio=enable_reference_audio,
    )

    # LLAMA Inference
    request = dict(
        device=decoder_model.device,
        max_new_tokens=max_new_tokens,
        text=text,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        temperature=temperature,
        compile=args.compile,
        iterative_prompt=chunk_length > 0,
        chunk_length=chunk_length,
        max_length=2048,
        prompt_tokens=prompt_tokens if enable_reference_audio else None,
        prompt_text=reference_text if enable_reference_audio else None,
    )

    response_queue = queue.Queue()
    llama_queue.put(
        GenerateRequest(
            request=request,
            response_queue=response_queue,
        )
    )

    if streaming:

        yield wav_chunk_header(), None, None

    segments = []

    while True:
        result: WrappedGenerateResponse = response_queue.get()
        if result.status == "error":
            yield None, None, build_html_error_message(result.response)
            break

        result: GenerateResponse = result.response
        if result.action == "next":
            break

        with autocast_exclude_mps(
            device_type=decoder_model.device.type, dtype=args.precision
        ):
            fake_audios = decode_vq_tokens(
                decoder_model=decoder_model,
                codes=result.codes,
            )

        fake_audios = fake_audios.float().cpu().numpy()
        segments.append(fake_audios)
        if streaming:
            yield (fake_audios * 32768).astype(np.int16).tobytes(), None, None

    if len(segments) == 0:
        return (
            None,
            None,
            build_html_error_message(
                i18n("No audio generated, please check the input text.")
            ),
        )

    # Concatenate audio segments
    audio = np.concatenate(segments, axis=0)
    return audio


def wav_chunk_header(sample_rate=44100, bit_depth=16, channels=1):
    buffer = io.BytesIO()

    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(channels)
        wav_file.setsampwidth(bit_depth // 8)
        wav_file.setframerate(sample_rate)

    wav_header_bytes = buffer.getvalue()
    buffer.close()
    return wav_header_bytes


def normalize_text(user_input, use_normalization):
    if use_normalization:
        return ChnNormedText(raw_text=user_input).normalize()
    else:
        return user_input

# 请求模型加载和初始化


def load_models():
    global decoder_model, llama_queue, args

    args = parse_args()
    args.precision = torch.half if args.half else torch.bfloat16

    logger.info("Loading Llama model...")
    llama_queue = launch_thread_safe_queue(
        checkpoint_path=args.llama_checkpoint_path,
        device=args.device,
        precision=args.precision,
        compile=args.compile,
    )
    logger.info("Llama model loaded, loading VQ-GAN model...")

    decoder_model = load_decoder_model(
        config_name=args.decoder_config_name,
        checkpoint_path=args.decoder_checkpoint_path,
        device=args.device,
    )

    logger.info("Decoder model loaded, warming up...")

    # Dry run to check if the model is loaded correctly and avoid the first-time latency
    list(
        inference(
            text="Hello, world!",
            enable_reference_audio=False,
            reference_audio=None,
            reference_text="",
            max_new_tokens=0,
            chunk_length=200,
            top_p=0.7,
            repetition_penalty=1.2,
            temperature=0.7,
        )
    )

    logger.info("Warming up done, ready to serve requests...")


def parse_args():
    parser = ArgumentParser()
    parser.add_argument(
        "--llama-checkpoint-path",
        type=Path,
        default="checkpoints/fish-speech-1.4",
    )
    parser.add_argument(
        "--decoder-checkpoint-path",
        type=Path,
        default="checkpoints/fish-speech-1.4/firefly-gan-vq-fsq-8x1024-21hz-generator.pth",
    )
    parser.add_argument("--decoder-config-name", type=str,
                        default="firefly_gan_vq")
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--half", action="store_true")
    parser.add_argument("--compile", action="store_true")
    parser.add_argument("--max-gradio-length", type=int, default=0)
    parser.add_argument("--theme", type=str, default="light")

    return parser.parse_args()

# 定义请求模型


class SynthesisRequest(BaseModel):
    text: str
    enable_reference_audio: bool = False
    reference_audio: bytes = None  # Base64 encoded audio file
    reference_text: str = ""
    max_new_tokens: int = 0
    chunk_length: int = 200
    top_p: float = 0.7
    repetition_penalty: float = 1.2
    temperature: float = 0.7
    seed: str = "0"
    use_normalization: bool = False


# 流式音频生成端点
@app.post("/stream_audio")
async def stream_audio(request: SynthesisRequest):
    text = normalize_text(request.text, request.use_normalization)

    # 如果有参考音频，保存到临时文件
    reference_audio_path = None
    if request.enable_reference_audio and request.reference_audio:
        reference_audio_path = "temp_reference_audio.wav"
        with open(reference_audio_path, "wb") as f:
            f.write(request.reference_audio)

    audio_generator = inference(
        text=text,
        enable_reference_audio=request.enable_reference_audio,
        reference_audio=reference_audio_path,
        reference_text=request.reference_text,
        max_new_tokens=request.max_new_tokens,
        chunk_length=request.chunk_length,
        top_p=request.top_p,
        repetition_penalty=request.repetition_penalty,
        temperature=request.temperature,
        seed=request.seed,
        streaming=True,
    )

    async def audio_stream():
        for chunk in audio_generator:
            audio_chunk, _, error_message = chunk
            if error_message:
                yield b""
                break
            if audio_chunk:
                yield audio_chunk

    headers = {
        "Content-Type": "audio/wav",
        "Transfer-Encoding": "chunked",
    }

    return StreamingResponse(audio_stream(), media_type="audio/wav", headers=headers)

if __name__ == "__main__":
    import uvicorn

    # 加载模型
    load_models()

    # 启动FastAPI应用
    uvicorn.run(app, host="0.0.0.0", port=8008)