Copy the files in the new repo for public release
This commit is contained in:
gabrieldemarmiesse
@@ -0,0 +1,254 @@
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli_app()
|
||||
Reference in New Issue
Block a user