Update app.py

app.py

@@ -2,40 +2,55 @@
 # 1) SETUP & IMPORTS
 # ===================================================================================
 from __future__ import annotations
+
 import os
 import sys
+import re
 import base64
 import struct
 import textwrap
 import requests
 import atexit
-from typing import List, Dict, Tuple, Generator
+from typing import List, Dict, Tuple, Generator, Any
 
 # --- Fast, safe defaults ---
 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")
-os.environ.setdefault("TORCHAUDIO_USE_FFMPEG", "0")  # prevent torchaudio/ffmpeg (torio) path
+os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")   # truly disable analytics
+os.environ.setdefault("TORCHAUDIO_USE_FFMPEG", "0")          # avoid torchaudio/ffmpeg linkage quirks
 
 # --- .env early (HF_TOKEN / SECRET_TOKEN) ---
 from dotenv import load_dotenv
 load_dotenv()
 
-# --- NumPy sanity with torch 2.2.x ---
+# --- NumPy sanity (Torch 2.2.x prefers NumPy 1.x) ---
 import numpy as _np
 if int(_np.__version__.split(".", 1)[0]) >= 2:
     raise RuntimeError(
-        f"Detected numpy=={_np.__version__}. Please pin numpy<2 (e.g., 1.26.4) for this Space."
+        f"Detected numpy=={_np.__version__}. Please ensure numpy<2 (e.g., 1.26.4)."
     )
 
-# --- Hugging Face Spaces & ZeroGPU (import BEFORE torch/diffusers) ---
+# --- Pandas compat shim (Gradio uses 'future.no_silent_downcasting') ---
+try:
+    import pandas as pd
+    from pandas._config.config import register_option
+    try:
+        pd.get_option("future.no_silent_downcasting")
+    except Exception:
+        # Register the option if missing so pd.option_context(...) won't crash
+        register_option("future.no_silent_downcasting", False, validator=None, doc="compat shim for Gradio")
+except Exception:
+    pd = None  # ok
+
+# --- Hugging Face Spaces & ZeroGPU (import BEFORE CUDA libs) ---
 try:
     import spaces  # Required for ZeroGPU on HF
 except Exception:
     class _SpacesShim:
         def GPU(self, *args, **kwargs):
-            def _wrap(fn): return fn
+            def _wrap(fn):
+                return fn
             return _wrap
     spaces = _SpacesShim()
 
@@ -47,7 +62,7 @@ import numpy as np
 from huggingface_hub import HfApi, hf_hub_download
 from llama_cpp import Llama
 
-# --- Audio decode via ffmpeg-python (no torchaudio.load) ---
+# --- Audio decoding (pure ffmpeg-python; no torchaudio) ---
 import ffmpeg
 
 # --- TTS Libraries ---
@@ -62,17 +77,30 @@ import langid
 import emoji
 import noisereduce as nr
 
+
 # ===================================================================================
 # 2) GLOBALS & HELPERS
 # ===================================================================================
 
-# NLTK data
-nltk.download("punkt", quiet=True)
+def _ensure_nltk() -> None:
+    # Newer NLTK splits data into 'punkt' and 'punkt_tab'
+    for pkg in ("punkt", "punkt_tab"):
+        try:
+            if pkg == "punkt":
+                nltk.data.find("tokenizers/punkt")
+            else:
+                nltk.data.find("tokenizers/punkt_tab")
+        except LookupError:
+            nltk.download(pkg, quiet=True)
 
-# Cached models & latents
+_ensure_nltk()
+
+# Models & caches
 tts_model: Xtts | None = None
 llm_model: Llama | None = None
-voice_latents: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+# Store latents as NumPy on CPU for portability; convert to device at inference time
+voice_latents: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
 
 # Config
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -82,9 +110,6 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "secret")
 SENTENCE_SPLIT_LENGTH = 250
 LLM_STOP_WORDS = ["</s>", "<|user|>", "/s>"]
 
-# IMPORTANT: With ZeroGPU, DO NOT use CUDA at startup even if torch sees it.
-USE_STARTUP_CUDA = os.getenv("USE_STARTUP_CUDA", "false").lower() == "true"
-
 # System prompts and roles
 default_system_message = (
     "You're a storyteller crafting a short tale for young listeners. Keep sentences short and simple. "
@@ -98,7 +123,25 @@ ROLE_PROMPTS["Pirate"] = (
     "Keep answers short, as if in a real conversation. Only provide the words AI Beard would speak."
 )
 
-# ---------- small utils ----------
+
+# ---------- tiny utilities ----------
+def _model_device(m: torch.nn.Module) -> torch.device:
+    try:
+        return next(m.parameters()).device
+    except StopIteration:
+        return torch.device("cpu")
+
+def _to_device_float_tensor(x: Any, device: torch.device) -> torch.Tensor:
+    if isinstance(x, np.ndarray):
+        return torch.from_numpy(x).float().to(device)
+    if torch.is_tensor(x):
+        return x.to(device, dtype=torch.float32)
+    return torch.as_tensor(x, dtype=torch.float32, device=device)
+
+def _latents_for_device(latents: Tuple[Any, Any], device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+    gpt_cond, spk = latents
+    return _to_device_float_tensor(gpt_cond, device), _to_device_float_tensor(spk, device)
+
 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
@@ -114,14 +157,19 @@ def pcm_to_wav(pcm_data: bytes, sample_rate: int = 24000, channels: int = 1, bit
     return header + pcm_data
 
 def split_sentences(text: str, max_len: int) -> List[str]:
-    sentences = nltk.sent_tokenize(text)
-    out: List[str] = []
-    for s in sentences:
-        if len(s) > max_len:
-            out.extend(textwrap.wrap(s, max_len, break_long_words=True))
+    # Try NLTK; if it fails for any reason, fallback to a simple regex splitter.
+    try:
+        sentences = nltk.sent_tokenize(text)
+    except Exception:
+        sentences = re.split(r"(?<=[\.\!\?])\s+", text)
+    chunks: List[str] = []
+    for sent in sentences:
+        if len(sent) > max_len:
+            chunks.extend(textwrap.wrap(sent, max_len, break_long_words=True))
         else:
-            out.append(s)
-    return out
+            if sent:
+                chunks.append(sent)
+    return chunks
 
 def format_prompt_zephyr(message: str, history: List[Tuple[str, str | None]], system_message: str) -> str:
     prompt = f"<|system|>\n{system_message}</s>"
@@ -131,6 +179,7 @@ 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 (mono via ffmpeg) ----------
 def _decode_audio_ffmpeg_to_mono(path: str, target_sr: int) -> np.ndarray:
     """
@@ -147,17 +196,23 @@ def _decode_audio_ffmpeg_to_mono(path: str, target_sr: int) -> np.ndarray:
         pcm = np.frombuffer(out, dtype=np.int16)
         if pcm.size == 0:
             raise RuntimeError("ffmpeg produced empty audio.")
-        return (pcm.astype(np.float32) / 32767.0)
+        wav = (pcm.astype(np.float32) / 32767.0)
+        return wav
     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.load() ----------
+
+# ---------- monkey-patch XTTS internal loader to avoid torchaudio/torio ----------
 def _patched_load_audio(audiopath: str, load_sr: int):
     """
-    Expected by XTTS: return torch.FloatTensor [1, samples] normalized to [-1, 1], resampled to load_sr.
+    Match XTTS' expected return type:
+    - returns a torch.FloatTensor shaped [1, samples], normalized to [-1, 1],
+      already resampled to `load_sr`.
+    - DO NOT return (audio, sr) tuple.
     """
     wav = _decode_audio_ffmpeg_to_mono(audiopath, target_sr=load_sr)
-    audio = torch.from_numpy(wav).float().unsqueeze(0)  # [1, N] on CPU
+    import torch as _torch  # local import to avoid circularities
+    audio = _torch.from_numpy(wav).float().unsqueeze(0)  # [1, N] on CPU
     return audio
 
 xtts_module.load_audio = _patched_load_audio
@@ -167,12 +222,14 @@ try:
 except Exception:
     pass
 
+
 def _coqui_cache_dir() -> str:
-    # Coqui cache default on Linux
+    # 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)
+# 3) PRECACHE & MODEL LOADERS (CPU at startup to avoid ZeroGPU issues)
 # ===================================================================================
 
 def precache_assets() -> None:
@@ -205,8 +262,9 @@ def precache_assets() -> None:
     except Exception as e:
         print(f"Warning: GGUF pre-cache error: {e}")
 
-def _load_xtts(device: str) -> Xtts:
-    """Load XTTS from the local cache. Always CPU at startup for ZeroGPU compatibility."""
+
+def _load_xtts(device: str = "cpu") -> Xtts:
+    """Load XTTS from the local cache. Keep CPU at startup to avoid ZeroGPU device mixups."""
     print(f"Loading Coqui XTTS V2 model on {device.upper()}...")
     model_name = "tts_models/multilingual/multi-dataset/xtts_v2"
     ModelManager().download_model(model_name)  # idempotent
@@ -225,16 +283,20 @@ def _load_xtts(device: str) -> Xtts:
     print("XTTS model loaded.")
     return model
 
-def _load_llama_cpu_only() -> Llama:
-    """Load Llama (Zephyr GGUF) on CPU only (ZeroGPU friendly)."""
-    print("Loading LLM (Zephyr GGUF) on CPU...")
+
+def _load_llama() -> Llama:
+    """
+    Load Llama (Zephyr GGUF).
+    Keep simple & robust: default to CPU (works everywhere).
+    """
+    print("Loading LLM (Zephyr GGUF)...")
     zephyr_model_path = hf_hub_download(
         repo_id="TheBloke/zephyr-7B-beta-GGUF",
         filename="zephyr-7b-beta.Q5_K_M.gguf"
     )
     llm = Llama(
         model_path=zephyr_model_path,
-        n_gpu_layers=0,   # never touch CUDA at startup
+        n_gpu_layers=0,     # CPU-only for reliability across Spaces/ZeroGPU
         n_ctx=4096,
         n_batch=512,
         verbose=False
@@ -242,27 +304,27 @@ def _load_llama_cpu_only() -> Llama:
     print("LLM loaded (CPU).")
     return llm
 
+
 def init_models_and_latents() -> None:
     """
-    Preload TTS and LLM on CPU and compute voice latents on CPU.
-    This avoids any CUDA tensors outside the @spaces.GPU window.
+    Preload models on CPU and compute voice latents on CPU.
+    This avoids ZeroGPU's "mixed device" errors from torchaudio-based resampling.
     """
     global tts_model, llm_model, voice_latents
 
-    # Always CPU here (ZeroGPU rule)
-    target_device = "cpu"
-
     if tts_model is None:
-        tts_model = _load_xtts(device=target_device)
-    else:
-        tts_model.to("cpu")
+        tts_model = _load_xtts(device="cpu")  # always CPU at startup
 
     if llm_model is None:
-        llm_model = _load_llama_cpu_only()
+        llm_model = _load_llama()
 
-    # Pre-compute latents once on CPU (uses our ffmpeg loader)
     if not voice_latents:
         print("Computing voice conditioning latents (CPU)...")
+        # Ensure the TTS model is on CPU while computing latents
+        orig_dev = _model_device(tts_model)
+        if orig_dev.type != "cpu":
+            tts_model.to("cpu")
+
         with torch.no_grad():
             for role, filename in [
                 ("Cloée", "cloee-1.wav"),
@@ -271,11 +333,21 @@ def init_models_and_latents() -> None:
                 ("Thera", "thera-1.wav"),
             ]:
                 path = os.path.join("voices", filename)
-                # Returns torch tensors; keep them on CPU
-                voice_latents[role] = tts_model.get_conditioning_latents(
+                gpt_lat, spk_emb = tts_model.get_conditioning_latents(
                     audio_path=path, gpt_cond_len=30, max_ref_length=60
                 )
-        print("Voice latents ready (CPU).")
+                # Store as NumPy on CPU; convert to device on demand later
+                voice_latents[role] = (
+                    gpt_lat.detach().cpu().numpy(),
+                    spk_emb.detach().cpu().numpy(),
+                )
+
+        # Return model to original device (keep CPU at startup for safety)
+        if orig_dev.type != "cpu":
+            tts_model.to(orig_dev)
+
+        print("Voice latents ready.")
+
 
 # Ensure we close Llama cleanly to avoid __del__ issues at interpreter shutdown
 def _close_llm():
@@ -287,6 +359,7 @@ def _close_llm():
         pass
 atexit.register(_close_llm)
 
+
 # ===================================================================================
 # 4) INFERENCE HELPERS
 # ===================================================================================
@@ -294,17 +367,17 @@ atexit.register(_close_llm)
 def generate_text_stream(llm_instance: Llama, prompt: str,
                          history: List[Tuple[str, str | None]],
                          system_message_text: str) -> Generator[str, None, None]:
-    formatted = format_prompt_zephyr(prompt, history, system_message_text)
+    formatted_prompt = format_prompt_zephyr(prompt, history, system_message_text)
     stream = llm_instance(
-        formatted,
+        formatted_prompt,
         temperature=0.7,
         max_tokens=512,
         top_p=0.95,
         stop=LLM_STOP_WORDS,
         stream=True
     )
-    for resp in stream:
-        ch = resp["choices"][0]["text"]
+    for response in stream:
+        ch = response["choices"][0]["text"]
         try:
             is_single_emoji = (len(ch) == 1 and emoji.is_emoji(ch))
         except Exception:
@@ -313,31 +386,29 @@ def generate_text_stream(llm_instance: Llama, prompt: str,
             continue
         yield ch
 
-def _latents_to_device(latents: Tuple[torch.Tensor, torch.Tensor], device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
-    g, s = latents
-    if isinstance(g, torch.Tensor):
-        g = g.to(device)
-    if isinstance(s, torch.Tensor):
-        s = s.to(device)
-    return g, s
 
 def generate_audio_stream(tts_instance: Xtts, text: str, language: str,
-                          latents: Tuple[torch.Tensor, torch.Tensor]) -> Generator[bytes, None, None]:
-    gpt_cond_latent, speaker_embedding = latents
+                          latents: Tuple[np.ndarray, np.ndarray]) -> Generator[bytes, None, None]:
+    # Convert stored CPU NumPy latents to tensors on the model's current device
+    device = _model_device(tts_instance)
+    gpt_cond_latent_t, speaker_embedding_t = _latents_for_device(latents, device)
+
     try:
         for chunk in tts_instance.inference_stream(
             text=text,
             language=language,
-            gpt_cond_latent=gpt_cond_latent,
-            speaker_embedding=speaker_embedding,
+            gpt_cond_latent=gpt_cond_latent_t,
+            speaker_embedding=speaker_embedding_t,
             temperature=0.85,
         ):
             if chunk is None:
                 continue
-            f32 = chunk.detach().cpu().numpy().squeeze().astype(np.float32)
-            f32 = np.clip(f32, -1.0, 1.0)
+            # chunk: torch.FloatTensor [N] or [1, N], float32 in [-1, 1]
+            f32 = chunk.detach().cpu().numpy().squeeze()
+            f32 = np.clip(f32, -1.0, 1.0).astype(np.float32)
             s16 = (f32 * 32767.0).astype(np.int16)
             yield s16.tobytes()
+
     except RuntimeError as e:
         print(f"Error during TTS inference: {e}")
         if "device-side assert" in str(e) and api:
@@ -347,34 +418,32 @@ def generate_audio_stream(tts_instance: Xtts, text: str, language: str,
             except Exception:
                 pass
 
+
 # ===================================================================================
-# 5) ZERO-GPU ENTRYPOINT (also works on native GPU)
+# 5) ZERO-GPU ENTRYPOINT (safe on native GPU as well)
 # ===================================================================================
 
-@spaces.GPU(duration=120)  # ZeroGPU allocates a GPU only for this function call
+@spaces.GPU(duration=120)  # GPU ops must occur inside this function when on ZeroGPU
 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 (CPU)
+    # Ensure models/latents exist (loaded on CPU)
     if tts_model is None or llm_model is None or not voice_latents:
         init_models_and_latents()
 
-    # If ZeroGPU granted CUDA for this call, move XTTS to CUDA; keep LLM on CPU.
+    # During the GPU window, move XTTS to CUDA if available; otherwise stay on CPU
     try:
         if torch.cuda.is_available():
             tts_model.to("cuda")
-            device = torch.device("cuda")
         else:
             tts_model.to("cpu")
-            device = torch.device("cpu")
     except Exception:
         tts_model.to("cpu")
-        device = torch.device("cpu")
 
-    # Generate story text (LLM on CPU)
+    # Generate story text (LLM kept CPU for simplicity & reliability)
     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])
@@ -391,13 +460,10 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
         if not any(c.isalnum() for c in sentence):
             continue
 
-        # Move cached latents to the same device as the model for this call
-        lat_dev = _latents_to_device(voice_latents[chatbot_role], device)
-
-        audio_chunks = generate_audio_stream(tts_model, sentence, lang, lat_dev)
+        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, CPU)
+        # Optional noise reduction (best-effort)
         try:
             data_s16 = np.frombuffer(pcm_data, dtype=np.int16)
             if data_s16.size > 0:
@@ -412,7 +478,7 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
         b64_wav = base64.b64encode(pcm_to_wav(final_pcm, sample_rate=24000, channels=1, bit_depth=16)).decode("utf-8")
         results.append({"text": sentence, "audio": b64_wav})
 
-    # Return XTTS to CPU to release GPU instantly
+    # Leave model on CPU after the ZeroGPU window
     try:
         tts_model.to("cpu")
     except Exception:
@@ -420,6 +486,7 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
 
     return results
 
+
 # ===================================================================================
 # 6) STARTUP: PRECACHE & UI
 # ===================================================================================
@@ -434,17 +501,21 @@ def build_ui() -> gr.Interface:
         ],
         outputs=gr.JSON(label="Story and Audio Output"),
         title="AI Storyteller with ZeroGPU",
-        description="Enter a prompt to generate a short story with voice narration using on-demand GPU.",
-        allow_flagging="never",
-        analytics_enabled=False,
+        description="Enter a prompt to generate a short story with voice narration. Uses GPU only within the generation call when available.",
+        flagging_mode="never",
+        # removed allow_flagging (deprecated) to silence warning
     )
 
 if __name__ == "__main__":
-    print("===== Startup: pre-cache assets and preload models =====")
-    print(f"Python: {sys.version.split()[0]} | Torch CUDA visible: {torch.cuda.is_available()} (will not use at startup)")
-    precache_assets()         # 1) download everything to disk
-    init_models_and_latents() # 2) load on CPU + compute voice latents on CPU
+    print("===== Startup: pre-cache assets and preload models (CPU) =====")
+    print(f"Python: {sys.version.split()[0]} | Torch CUDA available: {torch.cuda.is_available()}")
+    precache_assets()              # 1) download everything to disk
+    init_models_and_latents()      # 2) load models on CPU + compute voice latents on CPU
     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")),
+        ssr_mode=False,  # disable experimental SSR noise
+    )