Limit runtime admission to prevent OOM

Dockerfile

@@ -71,7 +71,11 @@ ENV HOME=/home/user \
     BACKEND_HOST=127.0.0.1 \
     BACKEND_PORT=8010 \
     AUDIO8_TTS_MEM_FRACTION_STATIC=0.60 \
-    AUDIO8_TTS_MAX_RUNNING_REQUESTS=8 \
+    AUDIO8_TTS_MAX_RUNNING_REQUESTS=2 \
+    AUDIO8_TTS_MAX_TEXT_UNITS=150 \
+    AUDIO8_TTS_MAX_RAW_TEXT_CHARS=1000 \
+    AUDIO8_TTS_MAX_NEW_TOKENS=1024 \
+    UI_MAX_CONCURRENCY=2 \
     AUDIO8_TTS_DISABLE_CUDA_GRAPH=0 \
     AUDIO8_TTS_ENABLE_TORCH_COMPILE=1 \
     AUDIO8_TTS_VERIFY_FAST_KV=1 \

app.py

@@ -4,6 +4,7 @@ import asyncio
 import os
 import tempfile
 import time
+import unicodedata
 from pathlib import Path
 from typing import Annotated
 
@@ -18,7 +19,10 @@ ROOT = Path(__file__).resolve().parent
 BACKEND_URL = os.getenv("BACKEND_URL", "http://127.0.0.1:8010")
 MAX_UPLOAD_BYTES = 15 * 1024 * 1024
 ALLOWED_AUDIO_SUFFIXES = {".wav", ".mp3", ".flac", ".m4a", ".ogg"}
-GENERATION_LIMIT = asyncio.Semaphore(int(os.getenv("UI_MAX_CONCURRENCY", "4")))
+MAX_TEXT_UNITS = int(os.getenv("AUDIO8_TTS_MAX_TEXT_UNITS", "150"))
+MAX_RAW_TEXT_CHARS = int(os.getenv("AUDIO8_TTS_MAX_RAW_TEXT_CHARS", "1000"))
+MAX_NEW_TOKENS = int(os.getenv("AUDIO8_TTS_MAX_NEW_TOKENS", "1024"))
+GENERATION_LIMIT = asyncio.Semaphore(int(os.getenv("UI_MAX_CONCURRENCY", "2")))
 
 ENGLISH_REFERENCE_TEXT = (
     "hello nice to meet you, what would you like to talk about todat"
@@ -107,7 +111,94 @@ async def _backend_status() -> tuple[bool, dict]:
         return False, {}
 
 
+def _normalize_speech_text(text: str) -> str:
+    cleaned: list[str] = []
+    for character in text:
+        if character.isspace():
+            cleaned.append(" ")
+        elif unicodedata.category(character) not in {"Cc", "Cf", "Cs", "Co", "Cn"}:
+            cleaned.append(character)
+    return " ".join("".join(cleaned).split())
+
+
+def _is_cjk_character(character: str) -> bool:
+    codepoint = ord(character)
+    return (
+        0x3400 <= codepoint <= 0x4DBF
+        or 0x4E00 <= codepoint <= 0x9FFF
+        or 0xF900 <= codepoint <= 0xFAFF
+        or 0x3040 <= codepoint <= 0x30FF
+        or 0xAC00 <= codepoint <= 0xD7AF
+    )
+
+
+def _count_speech_units(text: str) -> int:
+    units = 0
+    in_latin_word = False
+    for character in text:
+        if _is_cjk_character(character):
+            units += 1
+            in_latin_word = False
+            continue
+
+        name = unicodedata.name(character, "")
+        if character.isnumeric() or (character.isalpha() and "LATIN" in name):
+            if not in_latin_word:
+                units += 1
+            in_latin_word = True
+        elif character in {"'", "\u2019", "-"} and in_latin_word:
+            continue
+        else:
+            in_latin_word = False
+            if character.isalpha() or character.isnumeric():
+                units += 1
+    return units
+
+
+def _validate_speech_text(text: str) -> str:
+    text = _normalize_speech_text(text)
+    if len(text) > MAX_RAW_TEXT_CHARS:
+        raise HTTPException(status_code=400, detail="Speech text is too long")
+    if not text:
+        raise HTTPException(status_code=400, detail="Text must not be empty")
+    units = _count_speech_units(text)
+    if units == 0:
+        raise HTTPException(status_code=400, detail="Text must contain readable characters")
+    if units > MAX_TEXT_UNITS:
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                f"Text must be {MAX_TEXT_UNITS} Chinese characters or "
+                "English words or fewer"
+            ),
+        )
+    return text
+
+
+def _validate_speech_payload(payload: dict) -> dict:
+    text = payload.get("input")
+    if not isinstance(text, str):
+        raise HTTPException(status_code=400, detail="Text must not be empty")
+    text = _validate_speech_text(text)
+
+    max_new_tokens = payload.get("max_new_tokens", MAX_NEW_TOKENS)
+    if isinstance(max_new_tokens, bool):
+        raise HTTPException(status_code=400, detail="Max tokens must be an integer")
+    try:
+        max_new_tokens = int(max_new_tokens)
+    except (TypeError, ValueError) as exc:
+        raise HTTPException(status_code=400, detail="Max tokens must be an integer") from exc
+    if not 32 <= max_new_tokens <= MAX_NEW_TOKENS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Max tokens must be between 32 and {MAX_NEW_TOKENS}",
+        )
+
+    return {**payload, "input": text, "max_new_tokens": max_new_tokens}
+
+
 async def _generate(payload: dict) -> httpx.Response:
+    payload = _validate_speech_payload(payload)
     async with GENERATION_LIMIT:
         async with httpx.AsyncClient(timeout=httpx.Timeout(600.0)) as client:
             return await client.post(f"{BACKEND_URL}/v1/audio/speech", json=payload)
@@ -189,12 +280,8 @@ async def generate_speech(
     top_k: Annotated[int, Form()] = 50,
     max_new_tokens: Annotated[int, Form()] = 1024,
 ) -> Response:
-    text = " ".join(text.split())
+    text = _validate_speech_text(text)
     reference_text = " ".join(reference_text.split())
-    if not text:
-        raise HTTPException(status_code=400, detail="Text must not be empty")
-    if len(text) > 1000:
-        raise HTTPException(status_code=400, detail="Text must be 1000 characters or fewer")
     if not reference_text:
         raise HTTPException(status_code=400, detail="Reference transcript is required")
     if not 0 <= temperature <= 2:
@@ -203,8 +290,11 @@ async def generate_speech(
         raise HTTPException(status_code=400, detail="Top P must be between 0 and 1")
     if not 1 <= top_k <= 200:
         raise HTTPException(status_code=400, detail="Top K must be between 1 and 200")
-    if not 32 <= max_new_tokens <= 2048:
-        raise HTTPException(status_code=400, detail="Max tokens must be between 32 and 2048")
+    if not 32 <= max_new_tokens <= MAX_NEW_TOKENS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Max tokens must be between 32 and {MAX_NEW_TOKENS}",
+        )
 
     temporary_path: Path | None = None
     if reference_audio is not None and reference_audio.filename:
@@ -281,13 +371,19 @@ async def models_proxy() -> Response:
 
 @app.post("/v1/audio/speech")
 async def speech_proxy(request: Request) -> Response:
-    body = await request.body()
-    async with httpx.AsyncClient(timeout=httpx.Timeout(600.0)) as client:
-        response = await client.post(
-            f"{BACKEND_URL}/v1/audio/speech",
-            content=body,
-            headers={"Content-Type": request.headers.get("content-type", "application/json")},
-        )
+    try:
+        payload = await request.json()
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail="Request body must be valid JSON") from exc
+    if not isinstance(payload, dict):
+        raise HTTPException(status_code=400, detail="Request body must be a JSON object")
+
+    try:
+        response = await _generate(payload)
+    except httpx.ConnectError as exc:
+        raise HTTPException(status_code=503, detail="Model is still warming up") from exc
+    except httpx.TimeoutException as exc:
+        raise HTTPException(status_code=504, detail="Generation timed out") from exc
     forwarded_headers = {
         name: value
         for name, value in response.headers.items()

static/app.js

@@ -9,6 +9,11 @@ const state = {
 };
 
 const MAX_REFERENCE_SECONDS = 30;
+const MAX_SPEECH_UNITS = 150;
+const INVISIBLE_CHARACTERS = /[\u0000-\u001f\u007f-\u009f\u200b-\u200f\u202a-\u202e\u2060-\u206f\ufeff]/g;
+const CJK_CHARACTER = /[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/u;
+const LATIN_OR_NUMBER = /[\p{Script=Latin}\p{Number}]/u;
+const LETTER_OR_NUMBER = /[\p{Letter}\p{Number}]/u;
 
 const textPresets = {
   zh: "今天想和你分享一个好消息，Audio8 现在可以用更高效的方式生成自然流畅的语音。",
@@ -57,6 +62,44 @@ function formatDuration(seconds) {
   return seconds < 60 ? `${seconds.toFixed(1)}s` : `${Math.floor(seconds / 60)}m ${Math.round(seconds % 60)}s`;
 }
 
+function normalizeSpeechText(value, trim = true) {
+  const withoutInvisible = value.replace(INVISIBLE_CHARACTERS, (character) => (
+    /\s/.test(character) ? " " : ""
+  ));
+  const normalized = withoutInvisible.replace(/\s+/g, " ");
+  return trim ? normalized.trim() : normalized;
+}
+
+function countSpeechUnits(value) {
+  let units = 0;
+  let inLatinWord = false;
+  for (const character of value) {
+    if (CJK_CHARACTER.test(character)) {
+      units += 1;
+      inLatinWord = false;
+    } else if (LATIN_OR_NUMBER.test(character)) {
+      if (!inLatinWord) units += 1;
+      inLatinWord = true;
+    } else if (["'", "\u2019", "-"].includes(character) && inLatinWord) {
+      continue;
+    } else {
+      inLatinWord = false;
+      if (LETTER_OR_NUMBER.test(character)) units += 1;
+    }
+  }
+  return units;
+}
+
+function updateSpeechLimit() {
+  const units = countSpeechUnits(elements.speechText.value);
+  elements.characterCount.textContent = units;
+  elements.speechText.setAttribute("aria-invalid", String(units > MAX_SPEECH_UNITS));
+  elements.generateButton.disabled = (
+    !elements.consent.checked || units === 0 || units > MAX_SPEECH_UNITS
+  );
+  return units;
+}
+
 function readAudioDuration(file) {
   return new Promise((resolve, reject) => {
     const probe = new Audio();
@@ -278,6 +321,13 @@ function setGenerating(generating) {
 async function submitGeneration(event) {
   event.preventDefault();
   if (!elements.consent.checked) return;
+  elements.speechText.value = normalizeSpeechText(elements.speechText.value);
+  const speechUnits = updateSpeechLimit();
+  if (speechUnits === 0 || speechUnits > MAX_SPEECH_UNITS) {
+    showOutput("error");
+    elements.errorMessage.textContent = `Speech text must be ${MAX_SPEECH_UNITS} Chinese characters or English words or fewer.`;
+    return;
+  }
   if (state.source === "upload" && !elements.referenceAudio.files.length) {
     showOutput("error");
     elements.errorMessage.textContent = "Choose a reference audio file first.";
@@ -358,11 +408,15 @@ elements.referenceAudio.addEventListener("change", () => {
 });
 
 elements.speechText.addEventListener("input", () => {
-  elements.characterCount.textContent = elements.speechText.value.length;
+  const normalized = normalizeSpeechText(elements.speechText.value, false);
+  if (normalized !== elements.speechText.value) {
+    elements.speechText.value = normalized;
+  }
+  updateSpeechLimit();
 });
 
 elements.consent.addEventListener("change", () => {
-  elements.generateButton.disabled = !elements.consent.checked;
+  updateSpeechLimit();
 });
 
 elements.resultAudio.addEventListener("loadedmetadata", () => {
@@ -371,7 +425,7 @@ elements.resultAudio.addEventListener("loadedmetadata", () => {
 
 elements.form.addEventListener("submit", submitGeneration);
 
-elements.characterCount.textContent = elements.speechText.value.length;
+updateSpeechLimit();
 refreshIcons();
 loadExamples().catch((error) => {
   elements.errorMessage.textContent = error.message;

static/index.html

@@ -123,7 +123,7 @@
                 <span class="step-index">02</span>
                 <h2>Speech text</h2>
               </div>
-              <span class="character-count"><span id="characterCount">0</span>/1000</span>
+              <span class="character-count"><span id="characterCount">0</span>/150</span>
             </div>
             <label class="sr-only" for="speechText">Text to generate</label>
             <textarea
@@ -161,7 +161,7 @@
                 </label>
                 <label>
                   Max tokens
-                  <input name="max_new_tokens" type="number" min="32" max="2048" step="32" value="1024" />
+                  <input name="max_new_tokens" type="number" min="32" max="1024" step="32" value="1024" />
                 </label>
               </div>
             </details>