262 lines
8.5 KiB
Python
262 lines
8.5 KiB
Python
import io
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
from pathlib import Path
|
|
from queue import Queue
|
|
|
|
import typer
|
|
import uvicorn
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import FileResponse, StreamingResponse
|
|
from typing_extensions import Annotated
|
|
|
|
from pocket_tts.data.audio import stream_audio_chunks
|
|
from pocket_tts.default_parameters import (
|
|
DEFAULT_AUDIO_PROMPT,
|
|
DEFAULT_EOS_THRESHOLD,
|
|
DEFAULT_FRAMES_AFTER_EOS,
|
|
DEFAULT_LSD_DECODE_STEPS,
|
|
DEFAULT_NOISE_CLAMP,
|
|
DEFAULT_TEMPERATURE,
|
|
DEFAULT_VARIANT,
|
|
)
|
|
from pocket_tts.models.tts_model import TTSModel
|
|
from pocket_tts.utils.logging_utils import enable_logging
|
|
from pocket_tts.utils.utils import PREDEFINED_VOICES, size_of_dict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
cli_app = typer.Typer(
|
|
help="Kyutai Pocket TTS - Text-to-Speech generation tool", pretty_exceptions_show_locals=False
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------
|
|
# The pocket-tts server implementation
|
|
# ------------------------------------------------------
|
|
|
|
# Global model instance
|
|
tts_model = None
|
|
global_model_state = None
|
|
|
|
web_app = FastAPI(
|
|
title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0"
|
|
)
|
|
web_app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=[
|
|
"http://localhost:3000",
|
|
"https://pod1-10007.internal.kyutai.org",
|
|
"https://kyutai.org",
|
|
],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@web_app.get("/")
|
|
async def root():
|
|
"""Serve the frontend."""
|
|
static_path = Path(__file__).parent / "static" / "index.html"
|
|
return FileResponse(static_path)
|
|
|
|
|
|
@web_app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy"}
|
|
|
|
|
|
def write_to_queue(queue, text_to_generate, model_state):
|
|
"""Allows writing to the StreamingResponse as if it were a file."""
|
|
|
|
class FileLikeToQueue(io.IOBase):
|
|
def __init__(self, queue):
|
|
self.queue = queue
|
|
|
|
def write(self, data):
|
|
self.queue.put(data)
|
|
|
|
def flush(self):
|
|
pass
|
|
|
|
def close(self):
|
|
self.queue.put(None)
|
|
|
|
audio_chunks = tts_model.generate_audio_stream(
|
|
model_state=model_state, text_to_generate=text_to_generate
|
|
)
|
|
stream_audio_chunks(FileLikeToQueue(queue), audio_chunks, tts_model.config.mimi.sample_rate)
|
|
|
|
|
|
def generate_data_with_state(text_to_generate: str, model_state: dict):
|
|
queue = Queue()
|
|
|
|
# Run your function in a thread
|
|
thread = threading.Thread(target=write_to_queue, args=(queue, text_to_generate, model_state))
|
|
thread.start()
|
|
|
|
# Yield data as it becomes available
|
|
i = 0
|
|
while True:
|
|
data = queue.get()
|
|
if data is None:
|
|
break
|
|
i += 1
|
|
yield data
|
|
|
|
thread.join()
|
|
|
|
|
|
@web_app.post("/tts")
|
|
def text_to_speech(
|
|
text: str = Form(...),
|
|
voice_url: str | None = Form(None),
|
|
voice_wav: UploadFile | None = File(None),
|
|
):
|
|
"""
|
|
Generate speech from text using the pre-loaded voice prompt or a custom voice.
|
|
|
|
Args:
|
|
text: Text to convert to speech
|
|
voice_url: Optional voice URL (http://, https://, or hf://)
|
|
voice_wav: Optional uploaded voice file (mutually exclusive with voice_url)
|
|
"""
|
|
if not text.strip():
|
|
raise HTTPException(status_code=400, detail="Text cannot be empty")
|
|
|
|
if voice_url is not None and voice_wav is not None:
|
|
raise HTTPException(status_code=400, detail="Cannot provide both voice_url and voice_wav")
|
|
|
|
# Use the appropriate model state
|
|
if voice_url is not None:
|
|
if not (
|
|
voice_url.startswith("http://")
|
|
or voice_url.startswith("https://")
|
|
or voice_url.startswith("hf://")
|
|
or voice_url in PREDEFINED_VOICES
|
|
):
|
|
raise HTTPException(
|
|
status_code=400, detail="voice_url must start with http://, https://, or hf://"
|
|
)
|
|
model_state = tts_model._cached_get_state_for_audio_prompt(voice_url, truncate=True)
|
|
logging.warning("Using voice from URL: %s", voice_url)
|
|
elif voice_wav is not None:
|
|
# Use uploaded voice file
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
|
|
content = voice_wav.file.read()
|
|
temp_file.write(content)
|
|
temp_file.flush()
|
|
|
|
try:
|
|
model_state = tts_model.get_state_for_audio_prompt(
|
|
Path(temp_file.name), truncate=True
|
|
)
|
|
finally:
|
|
os.unlink(temp_file.name)
|
|
else:
|
|
# Use default global model state
|
|
model_state = global_model_state
|
|
|
|
return StreamingResponse(
|
|
generate_data_with_state(text, model_state),
|
|
media_type="audio/wav",
|
|
headers={
|
|
"Content-Disposition": "attachment; filename=generated_speech.wav",
|
|
"Transfer-Encoding": "chunked",
|
|
},
|
|
)
|
|
|
|
|
|
@cli_app.command()
|
|
def serve(
|
|
voice: Annotated[
|
|
str, typer.Option(help="Path to voice prompt audio file (voice to clone)")
|
|
] = DEFAULT_AUDIO_PROMPT,
|
|
host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost",
|
|
port: Annotated[int, typer.Option(help="Port to bind to")] = 8000,
|
|
reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False,
|
|
):
|
|
"""Start the FastAPI server."""
|
|
|
|
global tts_model, global_model_state
|
|
tts_model = TTSModel.load_model(DEFAULT_VARIANT)
|
|
|
|
# Pre-load the voice prompt
|
|
global_model_state = tts_model.get_state_for_audio_prompt(voice)
|
|
logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB")
|
|
|
|
uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload)
|
|
|
|
|
|
# ------------------------------------------------------
|
|
# The pocket-tts single generation CLI implementation
|
|
# ------------------------------------------------------
|
|
|
|
|
|
@cli_app.command()
|
|
def generate(
|
|
text: Annotated[
|
|
str, typer.Option(help="Text to generate")
|
|
] = "Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.",
|
|
voice: Annotated[
|
|
str, typer.Option(help="Path to audio conditioning file (voice to clone)")
|
|
] = DEFAULT_AUDIO_PROMPT,
|
|
quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False,
|
|
variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT,
|
|
lsd_decode_steps: Annotated[
|
|
int, typer.Option(help="Number of generation steps")
|
|
] = DEFAULT_LSD_DECODE_STEPS,
|
|
temperature: Annotated[
|
|
float, typer.Option(help="Temperature for generation")
|
|
] = DEFAULT_TEMPERATURE,
|
|
noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
|
|
eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
|
|
frames_after_eos: Annotated[
|
|
int, typer.Option(help="Number of frames to generate after EOS")
|
|
] = DEFAULT_FRAMES_AFTER_EOS,
|
|
output_path: Annotated[
|
|
str, typer.Option(help="Output path for generated audio")
|
|
] = "./tts_output.wav",
|
|
device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
|
|
):
|
|
"""Generate speech using Kyutai Pocket TTS."""
|
|
if "cuda" in device:
|
|
# Cuda graphs capturing does not play nice with multithreading.
|
|
os.environ["NO_CUDA_GRAPH"] = "1"
|
|
|
|
log_level = logging.ERROR if quiet else logging.INFO
|
|
with enable_logging("pocket_tts", log_level):
|
|
tts_model = TTSModel.load_model(
|
|
variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold
|
|
)
|
|
tts_model.to(device)
|
|
|
|
model_state_for_voice = tts_model.get_state_for_audio_prompt(voice)
|
|
# Stream audio generation directly to file or stdout
|
|
audio_chunks = tts_model.generate_audio_stream(
|
|
model_state=model_state_for_voice,
|
|
text_to_generate=text,
|
|
frames_after_eos=frames_after_eos,
|
|
)
|
|
|
|
stream_audio_chunks(output_path, audio_chunks, tts_model.config.mimi.sample_rate)
|
|
|
|
# Only print the result message if not writing to stdout
|
|
if output_path != "-":
|
|
logger.info("Results written in %s", output_path)
|
|
logger.info("-" * 20)
|
|
logger.info(
|
|
"If you want to try multiple voices and prompts quickly, try the `serve` command."
|
|
)
|
|
logger.info(
|
|
"If you like Kyutai projects, comment, like, subscribe at https://x.com/kyutai_labs"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli_app()
|