updt

app.py

@@ -3,7 +3,6 @@
 # ===================================================================================
 from __future__ import annotations
 import os
-import sys
 import base64
 import struct
 import textwrap
@@ -15,23 +14,31 @@ from typing import List, Dict, Tuple, Generator
 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 os.environ.setdefault("COQUI_TOS_AGREED", "1")
-os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")  # truly disable analytics
-os.environ.setdefault("TORCHAUDIO_USE_FFMPEG", "0")  # avoid torchaudio ffmpeg path entirely
+os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")
 
-# --- .env early (HF_TOKEN / SECRET_TOKEN) ---
+# --- Prefer torchaudio sox_io/soundfile backend (avoid FFmpeg/torio bug) ---
+try:
+    import torchaudio
+    _backend_set = False
+    for _cand in ("sox_io", "soundfile"):
+        try:
+            torchaudio.set_audio_backend(_cand)
+            _backend_set = True
+            break
+        except Exception:
+            pass
+    if not _backend_set:
+        os.environ["TORCHAUDIO_USE_FFMPEG"] = "0"
+except Exception:
+    torchaudio = None
+
+# --- Load .env early (HF_TOKEN / SECRET_TOKEN) ---
 from dotenv import load_dotenv
 load_dotenv()
 
-# --- NumPy sanity (Torch 2.2.2 expects NumPy 1.x in your stack) ---
-import numpy as _np
-if int(_np.__version__.split(".", 1)[0]) >= 2:
-    raise RuntimeError(
-        f"Detected numpy=={_np.__version__}. Please ensure numpy<2 (e.g., 1.26.4) for this Space."
-    )
-
-# --- Hugging Face Spaces & ZeroGPU (import BEFORE CUDA libs) ---
+# --- Hugging Face Spaces & ZeroGPU ---
 try:
-    import spaces  # Required for ZeroGPU on HF
+    import spaces
 except Exception:
     class _SpacesShim:
         def GPU(self, *args, **kwargs):
@@ -42,20 +49,17 @@ except Exception:
 
 import gradio as gr
 
-# --- Core ML & Data Libraries (after spaces import) ---
+# --- Core ML & Data Libraries ---
 import torch
 import numpy as np
 from huggingface_hub import HfApi, hf_hub_download
 from llama_cpp import Llama
 
-# --- Audio decoding (we'll use ffmpeg-python to avoid torchaudio/torio) ---
-import ffmpeg
-
 # --- TTS Libraries ---
 from TTS.tts.configs.xtts_config import XttsConfig
 from TTS.tts.models.xtts import Xtts
 from TTS.utils.manage import ModelManager
-import TTS.tts.models.xtts as xtts_module  # for monkey-patching load_audio
+from TTS.utils.generic_utils import get_user_data_dir
 
 # --- Text & Audio Processing ---
 import nltk
@@ -67,15 +71,12 @@ import noisereduce as nr
 # 2) GLOBALS & HELPERS
 # ===================================================================================
 
-# NLTK data
 nltk.download("punkt", quiet=True)
 
-# Cached models & latents
 tts_model: Xtts | None = None
 llm_model: Llama | None = None
 voice_latents: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
 
-# Config
 HF_TOKEN = os.environ.get("HF_TOKEN")
 api = HfApi(token=HF_TOKEN) if HF_TOKEN else None
 repo_id = "ruslanmv/ai-story-server"
@@ -83,7 +84,6 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "secret")
 SENTENCE_SPLIT_LENGTH = 250
 LLM_STOP_WORDS = ["</s>", "<|user|>", "/s>"]
 
-# System prompts and roles
 default_system_message = (
     "You're a storyteller crafting a short tale for young listeners. Keep sentences short and simple. "
     "Use narrative style only, without lists or complex words. Type numbers as words (e.g., 'ten')."
@@ -96,7 +96,6 @@ ROLE_PROMPTS["Pirate"] = (
     "Keep answers short, as if in a real conversation. Only provide the words AI Beard would speak."
 )
 
-# ---------- small utils ----------
 def pcm_to_wav(pcm_data: bytes, sample_rate: int = 24000, channels: int = 1, bit_depth: int = 16) -> bytes:
     if pcm_data.startswith(b"RIFF"):
         return pcm_data
@@ -129,44 +128,11 @@ def format_prompt_zephyr(message: str, history: List[Tuple[str, str | None]], sy
     prompt += f"<|user|>\n{message}</s><|assistant|>"
     return prompt
 
-# ---------- robust audio decode (24k mono via ffmpeg) ----------
-def _decode_audio_ffmpeg_to_24k_mono(path: str) -> Tuple[np.ndarray, int]:
-    """Return float32 waveform in [-1,1], 24 kHz mono."""
-    try:
-        out, _ = (
-            ffmpeg
-            .input(path)
-            .output("pipe:", format="s16le", acodec="pcm_s16le", ac=1, ar=24000)
-            .run(capture_stdout=True, capture_stderr=True, cmd="ffmpeg")
-        )
-        pcm = np.frombuffer(out, dtype=np.int16)
-        if pcm.size == 0:
-            raise RuntimeError("ffmpeg produced empty audio.")
-        wav = (pcm.astype(np.float32) / 32767.0).copy()
-        return wav, 24000
-    except ffmpeg.Error as e:
-        raise RuntimeError(f"ffmpeg decode failed: {e.stderr.decode(errors='ignore') if e.stderr else e}") from e
-
-# ---------- monkey-patch XTTS internal loader to avoid torchaudio/torio ----------
-def _patched_load_audio(audiopath: str, load_sr: int):
-    wav, sr = _decode_audio_ffmpeg_to_24k_mono(audiopath)
-    # XTTS expects (audio, sr) and will handle truncation/conditioning windows.
-    return wav, sr
-
-xtts_module.load_audio = _patched_load_audio  # <- critical fix
-
-# ---------- where Coqui caches models (avoid get_user_data_dir import) ----------
-def _coqui_cache_dir() -> str:
-    # Matches what TTS uses on Linux: ~/.local/share/tts
-    return os.path.join(os.path.expanduser("~"), ".local", "share", "tts")
-
 # ===================================================================================
 # 3) PRECACHE & MODEL LOADERS (RUN BEFORE FIRST INFERENCE)
 # ===================================================================================
 
 def precache_assets() -> None:
-    """Download voice WAVs, XTTS weights, and Zephyr GGUF to local cache before any inference."""
-    # Voices
     print("Pre-caching voice files...")
     file_names = ["cloee-1.wav", "julian-bedtime-style-1.wav", "pirate_by_coqui.wav", "thera-1.wav"]
     base_url = "https://raw.githubusercontent.com/ruslanmv/ai-story-server/main/voices/"
@@ -182,31 +148,27 @@ def precache_assets() -> None:
             except Exception as e:
                 print(f"Failed to download {name}: {e}")
 
-    # XTTS model files
     print("Pre-caching XTTS v2 model files...")
     ModelManager().download_model("tts_models/multilingual/multi-dataset/xtts_v2")
 
-    # LLM GGUF
     print("Pre-caching Zephyr GGUF...")
     try:
         hf_hub_download(
             repo_id="TheBloke/zephyr-7B-beta-GGUF",
-            filename="zephyr-7b-beta.Q5_K_M.gguf",
-            force_download=False
+            filename="zephyr-7b-beta.Q5_K_M.gguf"
         )
     except Exception as e:
         print(f"Warning: GGUF pre-cache error: {e}")
 
 def _load_xtts(device: str) -> Xtts:
-    """Load XTTS from the local cache. Use checkpoint_dir to avoid None path bugs."""
     print("Loading Coqui XTTS V2 model (CPU first)...")
     model_name = "tts_models/multilingual/multi-dataset/xtts_v2"
-    ModelManager().download_model(model_name)  # idempotent
-    model_dir = os.path.join(_coqui_cache_dir(), model_name.replace("/", "--"))
+    model_dir = os.path.join(get_user_data_dir("tts"), model_name.replace("/", "--"))
 
     cfg = XttsConfig()
     cfg.load_json(os.path.join(model_dir, "config.json"))
     model = Xtts.init_from_config(cfg)
+
     model.load_checkpoint(
         cfg,
         checkpoint_dir=model_dir,
@@ -218,7 +180,6 @@ def _load_xtts(device: str) -> Xtts:
     return model
 
 def _load_llama() -> Llama:
-    """Load Llama (Zephyr GGUF) on CPU so it's ready immediately."""
     print("Loading LLM (Zephyr GGUF) on CPU...")
     zephyr_model_path = hf_hub_download(
         repo_id="TheBloke/zephyr-7B-beta-GGUF",
@@ -226,7 +187,7 @@ def _load_llama() -> Llama:
     )
     llm = Llama(
         model_path=zephyr_model_path,
-        n_gpu_layers=0,   # CPU by default
+        n_gpu_layers=0,
         n_ctx=4096,
         n_batch=512,
         verbose=False
@@ -235,16 +196,14 @@ def _load_llama() -> Llama:
     return llm
 
 def init_models_and_latents() -> None:
-    """Preload TTS and LLM on CPU and compute voice latents once (using patched audio loader)."""
     global tts_model, llm_model, voice_latents
 
     if tts_model is None:
-        tts_model = _load_xtts(device="cpu")  # keep on CPU at startup
+        tts_model = _load_xtts(device="cpu")
 
     if llm_model is None:
         llm_model = _load_llama()
 
-    # Pre-compute latents once (CPU OK); uses patched loader (ffmpeg) under the hood
     if not voice_latents:
         print("Computing voice conditioning latents...")
         for role, filename in [
@@ -259,14 +218,13 @@ def init_models_and_latents() -> None:
             )
         print("Voice latents ready.")
 
-# Ensure we close Llama cleanly to avoid __del__ issues at interpreter shutdown
 def _close_llm():
     global llm_model
-    try:
-        if llm_model is not None:
+    if llm_model is not None:
+        try:
             llm_model.close()
-    except Exception:
-        pass
+        except Exception:
+            pass
 atexit.register(_close_llm)
 
 # ===================================================================================
@@ -311,8 +269,8 @@ def generate_audio_stream(tts_instance: Xtts, text: str, language: str,
     except RuntimeError as e:
         print(f"Error during TTS inference: {e}")
         if "device-side assert" in str(e) and api:
+            gr.Warning("Critical GPU error. Attempting to restart the Space...")
             try:
-                gr.Warning("Critical GPU error. Attempting to restart the Space...")
                 api.restart_space(repo_id=repo_id)
             except Exception:
                 pass
@@ -321,68 +279,57 @@ def generate_audio_stream(tts_instance: Xtts, text: str, language: str,
 # 5) ZERO-GPU ENTRYPOINT
 # ===================================================================================
 
-@spaces.GPU(duration=120)  # Request GPU for up to 120s; adjust as needed
+@spaces.GPU(duration=120)
 def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_role: str) -> List[Dict[str, str]]:
     if secret_token_input != SECRET_TOKEN:
         raise gr.Error("Invalid secret token provided.")
     if not input_text:
         return []
 
-    # Ensure models/latents exist
     if tts_model is None or llm_model is None or not voice_latents:
-        init_models_and_latents()
+        raise gr.Error("Models not initialized. Please restart the Space.")
 
-    # Move XTTS to CUDA for this call if GPU is available; otherwise stay on CPU
     try:
         if torch.cuda.is_available():
             tts_model.to("cuda")
         else:
             tts_model.to("cpu")
-    except Exception:
-        tts_model.to("cpu")
 
-    # Generate story text (LLM on CPU)
-    history: List[Tuple[str, str | None]] = [(input_text, None)]
-    full_story_text = "".join(
-        generate_text_stream(llm_model, history[-1][0], history[:-1], system_message_text=ROLE_PROMPTS[chatbot_role])
-    ).strip()
-    if not full_story_text:
-        return []
+        history: List[Tuple[str, str | None]] = [(input_text, None)]
+        full_story_text = "".join(
+            generate_text_stream(llm_model, history[-1][0], history[:-1], system_message_text=ROLE_PROMPTS[chatbot_role])
+        ).strip()
+        if not full_story_text:
+            return []
 
-    # Split into TTS-friendly sentences
-    sentences = split_sentences(full_story_text, SENTENCE_SPLIT_LENGTH)
-    lang = langid.classify(sentences[0])[0] if sentences else "en"
+        sentences = split_sentences(full_story_text, SENTENCE_SPLIT_LENGTH)
+        lang = langid.classify(sentences[0])[0] if sentences else "en"
 
-    results: List[Dict[str, str]] = []
-    for sentence in sentences:
-        if not any(c.isalnum() for c in sentence):
-            continue
+        results: List[Dict[str, str]] = []
+        for sentence in sentences:
+            if not any(c.isalnum() for c in sentence):
+                continue
 
-        audio_chunks = generate_audio_stream(tts_model, sentence, lang, voice_latents[chatbot_role])
-        pcm_data = b"".join(chunk for chunk in audio_chunks if chunk)
+            audio_chunks = generate_audio_stream(tts_model, sentence, lang, voice_latents[chatbot_role])
+            pcm_data = b"".join(chunk for chunk in audio_chunks if chunk)
 
-        # Optional noise reduction (best-effort)
-        try:
-            data_s16 = np.frombuffer(pcm_data, dtype=np.int16)
-            if data_s16.size > 0:
-                float_data = data_s16.astype(np.float32) / 32767.0
-                reduced = nr.reduce_noise(y=float_data, sr=24000)
-                final_pcm = (reduced * 32767).astype(np.int16).tobytes()
-            else:
+            try:
+                data_s16 = np.frombuffer(pcm_data, dtype=np.int16)
+                if data_s16.size > 0:
+                    float_data = data_s16.astype(np.float32) / 32767.0
+                    reduced = nr.reduce_noise(y=float_data, sr=24000)
+                    final_pcm = (reduced * 32767).astype(np.int16).tobytes()
+                else:
+                    final_pcm = pcm_data
+            except Exception:
                 final_pcm = pcm_data
-        except Exception:
-            final_pcm = pcm_data
 
-        b64_wav = base64.b64encode(pcm_to_wav(final_pcm)).decode("utf-8")
-        results.append({"text": sentence, "audio": b64_wav})
+            b64_wav = base64.b64encode(pcm_to_wav(final_pcm)).decode("utf-8")
+            results.append({"text": sentence, "audio": b64_wav})
 
-    # Release GPU immediately
-    try:
+        return results
+    finally:
         tts_model.to("cpu")
-    except Exception:
-        pass
-
-    return results
 
 # ===================================================================================
 # 6) STARTUP: PRECACHE & UI
@@ -404,9 +351,9 @@ def build_ui() -> gr.Interface:
 
 if __name__ == "__main__":
     print("===== Startup: pre-cache assets and preload models =====")
-    precache_assets()              # 1) download everything to disk
-    init_models_and_latents()      # 2) load models on CPU + compute voice latents (via patched loader)
+    precache_assets()
+    init_models_and_latents()
     print("Models and assets ready. Launching UI...")
 
     demo = build_ui()
-    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))

requirements.txt

@@ -9,7 +9,7 @@ requests
 numpy
 pandas==1.5.3
 
-# TTS
+# TTS (legacy)
 TTS @ git+https://github.com/coqui-ai/[email protected]
 pydantic==2.5.3
 
@@ -17,7 +17,6 @@ pydantic==2.5.3
 llama-cpp-python==0.2.79
 
 # Audio & Text
-soundfile
 noisereduce==3.0.3
 pydub
 langid
@@ -27,4 +26,4 @@ ffmpeg-python
 
 # Japanese Text (optional)
 mecab-python3==1.0.9
-unidic-lite==1.0.8
+unidic-lite==1.0.8