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)