first

server.py

@@ -0,0 +1,288 @@
+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)