289 lines
9.3 KiB
Python
289 lines
9.3 KiB
Python
import io
|
|
import re
|
|
import traceback
|
|
from typing import Literal
|
|
|
|
import numpy as np
|
|
import scipy.io.wavfile as wavfile
|
|
from fastapi import FastAPI, Form, HTTPException
|
|
from fastapi.openapi.docs import get_swagger_ui_html
|
|
from fastapi.responses import RedirectResponse, Response
|
|
from pocket_tts import TTSModel
|
|
from pocket_tts.models.tts_model import prepare_text_prompt, split_into_best_sentences
|
|
|
|
app = FastAPI(
|
|
title="Pocket TTS High-Quality API",
|
|
description="HTTP API for generating MP3 or WAV speech audio with Pocket TTS.",
|
|
version="1.0.0",
|
|
docs_url=None,
|
|
redoc_url="/redoc",
|
|
openapi_url="/openapi.json",
|
|
)
|
|
|
|
SWAGGER_UI_VERSION = "5.32.10"
|
|
SWAGGER_UI_JS_URL = (
|
|
f"https://cdn.jsdelivr.net/npm/swagger-ui-dist@{SWAGGER_UI_VERSION}/swagger-ui-bundle.js"
|
|
)
|
|
SWAGGER_UI_CSS_URL = (
|
|
f"https://cdn.jsdelivr.net/npm/swagger-ui-dist@{SWAGGER_UI_VERSION}/swagger-ui.css"
|
|
)
|
|
SAFE_MAX_TOKENS_PER_CHUNK = 35
|
|
JOIN_SILENCE_MS = 100
|
|
MIN_TOKENS_PER_CHUNK = 12
|
|
|
|
print("Caricamento del modello Pocket TTS in memoria...")
|
|
model = TTSModel.load_model(
|
|
language="spanish_24l", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
|
|
)
|
|
last_generated_audio: dict[str, bytes | str] | None = None
|
|
|
|
|
|
@app.get("/", include_in_schema=False)
|
|
async def swagger_ui_redirect():
|
|
return RedirectResponse(url="/docs")
|
|
|
|
|
|
@app.get("/docs", include_in_schema=False)
|
|
async def custom_swagger_ui():
|
|
return get_swagger_ui_html(
|
|
openapi_url=app.openapi_url,
|
|
title=f"{app.title} - Swagger UI",
|
|
swagger_js_url=SWAGGER_UI_JS_URL,
|
|
swagger_css_url=SWAGGER_UI_CSS_URL,
|
|
swagger_ui_parameters={
|
|
"displayRequestDuration": True,
|
|
},
|
|
)
|
|
|
|
|
|
def build_audio_response(audio_bytes: bytes, media_type: str) -> Response:
|
|
return Response(
|
|
content=audio_bytes,
|
|
media_type=media_type,
|
|
headers={"Cache-Control": "no-store"},
|
|
)
|
|
|
|
|
|
def normalize_tts_text(text: str) -> str:
|
|
def decode_unicode_match(match: re.Match[str]) -> str:
|
|
return match.group(0).encode("ascii").decode("unicode_escape")
|
|
|
|
if "\\u" in text:
|
|
text = re.sub(r"\\u[0-9a-fA-F]{4}", decode_unicode_match, text)
|
|
|
|
text = text.replace("\\n", "\n").replace("\\t", " ")
|
|
|
|
# Keep heading content while removing markdown wrappers.
|
|
text = re.sub(r"\*\*\s*\[([^\]]+)\]\s*\*\*", r"\n\1.\n", text)
|
|
text = re.sub(r"\[([^\]]+)\]", r"\1", text)
|
|
text = text.replace("**", "").replace("__", "")
|
|
|
|
# Remove list/heading markdown markers that harm pronunciation.
|
|
text = re.sub(r"^[#>*-]+\s*", "", text, flags=re.MULTILINE)
|
|
|
|
# Normalize whitespace while preserving sentence boundaries.
|
|
text = re.sub(r"\r\n?", "\n", text)
|
|
text = re.sub(r"\n{2,}", "\n\n", text)
|
|
text = re.sub(r"[ \t]+", " ", text)
|
|
text = re.sub(r" *\n *", "\n", text)
|
|
|
|
return text.strip()
|
|
|
|
|
|
def tensor_to_numpy(audio_tensor) -> np.ndarray:
|
|
if hasattr(audio_tensor, "detach"):
|
|
audio_numpy = audio_tensor.detach().cpu().numpy()
|
|
else:
|
|
audio_numpy = np.array(audio_tensor)
|
|
|
|
if audio_numpy.ndim > 1:
|
|
audio_numpy = audio_numpy.squeeze()
|
|
|
|
return np.asarray(audio_numpy, dtype=np.float32)
|
|
|
|
|
|
def count_model_tokens(text: str) -> int:
|
|
tokenizer = model.flow_lm.conditioner.tokenizer
|
|
prepared_text, _ = prepare_text_prompt(
|
|
text,
|
|
model.pad_with_spaces_for_short_inputs,
|
|
model.remove_semicolons,
|
|
)
|
|
return len(tokenizer(prepared_text.strip()).tokens[0].tolist())
|
|
|
|
|
|
def split_text_into_safe_chunks(text: str, max_tokens: int) -> list[str]:
|
|
tokenizer = model.flow_lm.conditioner.tokenizer
|
|
chunks = split_into_best_sentences(
|
|
tokenizer,
|
|
text,
|
|
max_tokens,
|
|
model.pad_with_spaces_for_short_inputs,
|
|
remove_semicolons=model.remove_semicolons,
|
|
)
|
|
|
|
refined_chunks: list[str] = []
|
|
for chunk in chunks:
|
|
chunk = chunk.strip()
|
|
token_count = count_model_tokens(chunk)
|
|
if token_count <= max_tokens:
|
|
refined_chunks.append(chunk)
|
|
continue
|
|
|
|
words = chunk.split()
|
|
current_words: list[str] = []
|
|
for word in words:
|
|
candidate_words = current_words + [word]
|
|
candidate_text = " ".join(candidate_words)
|
|
candidate_tokens = count_model_tokens(candidate_text)
|
|
if current_words and candidate_tokens > max_tokens:
|
|
refined_chunks.append(" ".join(current_words))
|
|
current_words = [word]
|
|
else:
|
|
current_words = candidate_words
|
|
|
|
if current_words:
|
|
refined_chunks.append(" ".join(current_words))
|
|
|
|
merged_chunks: list[str] = []
|
|
for chunk in refined_chunks:
|
|
chunk = chunk.strip()
|
|
if not chunk:
|
|
continue
|
|
|
|
chunk_tokens = count_model_tokens(chunk)
|
|
if not merged_chunks:
|
|
merged_chunks.append(chunk)
|
|
continue
|
|
|
|
previous_chunk = merged_chunks[-1]
|
|
previous_tokens = count_model_tokens(previous_chunk)
|
|
combined = f"{previous_chunk} {chunk}".strip()
|
|
combined_tokens = count_model_tokens(combined)
|
|
|
|
# Avoid tiny fragments when the merge still fits comfortably.
|
|
if chunk_tokens < MIN_TOKENS_PER_CHUNK and combined_tokens <= max_tokens:
|
|
merged_chunks[-1] = combined
|
|
continue
|
|
|
|
# Also merge if the previous chunk is tiny and the combined text still fits.
|
|
if previous_tokens < MIN_TOKENS_PER_CHUNK and combined_tokens <= max_tokens:
|
|
merged_chunks[-1] = combined
|
|
continue
|
|
|
|
merged_chunks.append(chunk)
|
|
|
|
for chunk in merged_chunks:
|
|
token_count = count_model_tokens(chunk)
|
|
if token_count > max_tokens:
|
|
raise ValueError(
|
|
f"Chunk exceeds hard token limit ({token_count} > {max_tokens}): {chunk[:120]!r}"
|
|
)
|
|
|
|
return merged_chunks
|
|
|
|
|
|
def generate_audio_from_chunks(voice_state: dict, text: str, sample_rate: int) -> np.ndarray:
|
|
chunks = split_text_into_safe_chunks(text, SAFE_MAX_TOKENS_PER_CHUNK)
|
|
audio_segments: list[np.ndarray] = []
|
|
silence = np.zeros(int(sample_rate * JOIN_SILENCE_MS / 1000), dtype=np.float32)
|
|
|
|
for index, chunk in enumerate(chunks):
|
|
audio_tensor = model.generate_audio(
|
|
voice_state,
|
|
chunk,
|
|
max_tokens=SAFE_MAX_TOKENS_PER_CHUNK,
|
|
copy_state=True,
|
|
)
|
|
audio_segments.append(tensor_to_numpy(audio_tensor))
|
|
if index < len(chunks) - 1:
|
|
audio_segments.append(silence.copy())
|
|
|
|
if not audio_segments:
|
|
return np.zeros(0, dtype=np.float32)
|
|
|
|
return np.concatenate(audio_segments)
|
|
|
|
|
|
@app.get("/v1/audio/speech", include_in_schema=False)
|
|
@app.get("/tts", include_in_schema=False)
|
|
async def get_last_generated_audio():
|
|
if last_generated_audio is None:
|
|
raise HTTPException(status_code=404, detail="No generated audio available yet")
|
|
return build_audio_response(
|
|
audio_bytes=last_generated_audio["audio_bytes"],
|
|
media_type=last_generated_audio["media_type"],
|
|
)
|
|
|
|
|
|
@app.post("/v1/audio/speech")
|
|
@app.post(
|
|
"/tts",
|
|
tags=["tts"],
|
|
summary="Generate speech audio",
|
|
description=(
|
|
"Generate speech from text and return the result as an MP3 file by default, "
|
|
"or as WAV when `format=wav`. "
|
|
"The same handler is available on `/v1/audio/speech` and `/tts`."
|
|
),
|
|
response_class=Response,
|
|
responses={
|
|
200: {
|
|
"description": "Generated MP3 or WAV audio",
|
|
"content": {"audio/mpeg": {}, "audio/wav": {}},
|
|
},
|
|
500: {"description": "Audio generation failed"},
|
|
},
|
|
)
|
|
async def generate_tts(
|
|
text: str = Form(...),
|
|
voice_url: str = Form("lola"),
|
|
format: Literal["mp3", "wav"] = Form("mp3"),
|
|
):
|
|
global last_generated_audio
|
|
try:
|
|
normalized_text = normalize_tts_text(text)
|
|
# 1. Recupera lo stato del prompt vocale (dict)
|
|
voice_state = model.get_state_for_audio_prompt(voice_url)
|
|
sample_rate = model.sample_rate if hasattr(model, "sample_rate") else 24000
|
|
audio_numpy = generate_audio_from_chunks(voice_state, normalized_text, sample_rate)
|
|
|
|
buffer = io.BytesIO()
|
|
|
|
if format == "wav":
|
|
wavfile.write(buffer, sample_rate, audio_numpy)
|
|
media_type = "audio/wav"
|
|
else:
|
|
try:
|
|
import soundfile as sf
|
|
except ImportError as exc:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=(
|
|
"MP3 output requires the optional `soundfile` dependency. "
|
|
"Install it with `uv add soundfile`."
|
|
),
|
|
) from exc
|
|
|
|
sf.write(buffer, audio_numpy, sample_rate, format="MP3")
|
|
media_type = "audio/mpeg"
|
|
|
|
buffer.seek(0)
|
|
audio_bytes = buffer.read()
|
|
|
|
last_generated_audio = {
|
|
"audio_bytes": audio_bytes,
|
|
"media_type": media_type,
|
|
}
|
|
return build_audio_response(audio_bytes=audio_bytes, media_type=media_type)
|
|
|
|
except Exception as e:
|
|
print("\n--- ERRORE DURANTE LA GENERAZIONE ---")
|
|
traceback.print_exc()
|
|
print("-------------------------------------\n")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8009)
|