Update app.py

app.py

@@ -2,55 +2,70 @@
 # 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, Any
+from typing import List, Dict, Tuple, Generator
 
 # --- 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")   # truly disable analytics
-os.environ.setdefault("TORCHAUDIO_USE_FFMPEG", "0")          # avoid torchaudio/ffmpeg linkage quirks
+os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")
+os.environ.setdefault("TORCHAUDIO_USE_FFMPEG", "0")  # prevent torchaudio/ffmpeg (torio) path
 
 # --- .env early (HF_TOKEN / SECRET_TOKEN) ---
 from dotenv import load_dotenv
 load_dotenv()
 
-# --- NumPy sanity (Torch 2.2.x prefers NumPy 1.x) ---
+# --- NumPy sanity with torch 2.2.x ---
 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)."
+        f"Detected numpy=={_np.__version__}. Please pin numpy<2 (e.g., 1.26.4) for this Space."
     )
 
-# --- Pandas compat shim (Gradio uses 'future.no_silent_downcasting') ---
+# --- Panda shim for Gradio on pandas<2.2 (no 'future.no_silent_downcasting') ---
 try:
     import pandas as pd
-    from pandas._config.config import register_option
     try:
-        pd.get_option("future.no_silent_downcasting")
+        with pd.option_context("future.no_silent_downcasting", True):
+            pass
     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")
+        from contextlib import contextmanager
+        _orig_option_context = pd.option_context
+
+        @contextmanager
+        def _patched_option_context(*args, **kwargs):
+            # filter out unsupported option pairs
+            filtered = []
+            i = 0
+            while i < len(args):
+                key = args[i]
+                val = args[i + 1] if i + 1 < len(args) else None
+                if key == "future.no_silent_downcasting":
+                    i += 2
+                    continue
+                filtered.extend([key, val])
+                i += 2
+            with _orig_option_context(*filtered, **kwargs):
+                yield
+
+        pd.option_context = _patched_option_context  # type: ignore[attr-defined]
 except Exception:
-    pd = None  # ok
+    pd = None  # noqa: N816
 
-# --- Hugging Face Spaces & ZeroGPU (import BEFORE CUDA libs) ---
+# --- Hugging Face Spaces & ZeroGPU (import BEFORE torch/diffusers) ---
 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()
 
@@ -62,7 +77,7 @@ import numpy as np
 from huggingface_hub import HfApi, hf_hub_download
 from llama_cpp import Llama
 
-# --- Audio decoding (pure ffmpeg-python; no torchaudio) ---
+# --- Audio decode via ffmpeg-python (no torchaudio.load) ---
 import ffmpeg
 
 # --- TTS Libraries ---
@@ -77,30 +92,32 @@ import langid
 import emoji
 import noisereduce as nr
 
-
 # ===================================================================================
 # 2) GLOBALS & HELPERS
 # ===================================================================================
 
-def _ensure_nltk() -> None:
-    # Newer NLTK splits data into 'punkt' and 'punkt_tab'
-    for pkg in ("punkt", "punkt_tab"):
+# Ensure NLTK resources (both 'punkt' and new 'punkt_tab' on newer NLTK)
+def _ensure_nltk():
+    try:
+        nltk.data.find("tokenizers/punkt")
+    except LookupError:
+        nltk.download("punkt", quiet=True)
+    try:
+        nltk.data.find("tokenizers/punkt_tab/english")
+    except LookupError:
         try:
-            if pkg == "punkt":
-                nltk.data.find("tokenizers/punkt")
-            else:
-                nltk.data.find("tokenizers/punkt_tab")
-        except LookupError:
-            nltk.download(pkg, quiet=True)
+            nltk.download("punkt_tab", quiet=True)
+        except Exception:
+            # fallback: downloading 'punkt' already satisfies older versions
+            pass
 
 _ensure_nltk()
 
-# Models & caches
+# Cached models & latents
 tts_model: Xtts | None = None
 llm_model: Llama | None = None
-
-# Store latents as NumPy on CPU for portability; convert to device at inference time
-voice_latents: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
+# store as torch.Tensors (CPU at startup)
+voice_latents: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
 
 # Config
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -110,7 +127,10 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "secret")
 SENTENCE_SPLIT_LENGTH = 250
 LLM_STOP_WORDS = ["</s>", "<|user|>", "/s>"]
 
-# System prompts and roles
+# 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"
+
+# Roles & prompts
 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')."
@@ -123,25 +143,7 @@ ROLE_PROMPTS["Pirate"] = (
     "Keep answers short, as if in a real conversation. Only provide the words AI Beard would speak."
 )
 
-
-# ---------- 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)
-
+# ---------- 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
@@ -157,19 +159,14 @@ 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]:
-    # 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))
+    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))
         else:
-            if sent:
-                chunks.append(sent)
-    return chunks
+            out.append(s)
+    return out
 
 def format_prompt_zephyr(message: str, history: List[Tuple[str, str | None]], system_message: str) -> str:
     prompt = f"<|system|>\n{system_message}</s>"
@@ -179,7 +176,6 @@ 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:
     """
@@ -196,23 +192,17 @@ 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.")
-        wav = (pcm.astype(np.float32) / 32767.0)
-        return wav
+        return (pcm.astype(np.float32) / 32767.0)
     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 ----------
+# ---------- monkey-patch XTTS internal loader to avoid torchaudio.load() ----------
 def _patched_load_audio(audiopath: str, load_sr: int):
     """
-    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.
+    Expected by XTTS: return torch.FloatTensor [1, samples] normalized to [-1, 1], resampled to load_sr.
     """
     wav = _decode_audio_ffmpeg_to_mono(audiopath, target_sr=load_sr)
-    import torch as _torch  # local import to avoid circularities
-    audio = _torch.from_numpy(wav).float().unsqueeze(0)  # [1, N] on CPU
+    audio = torch.from_numpy(wav).float().unsqueeze(0)  # [1, N] on CPU
     return audio
 
 xtts_module.load_audio = _patched_load_audio
@@ -222,14 +212,12 @@ try:
 except Exception:
     pass
 
-
 def _coqui_cache_dir() -> str:
-    # Matches what TTS uses on Linux: ~/.local/share/tts
+    # Coqui cache default on Linux
     return os.path.join(os.path.expanduser("~"), ".local", "share", "tts")
 
-
 # ===================================================================================
-# 3) PRECACHE & MODEL LOADERS (CPU at startup to avoid ZeroGPU issues)
+# 3) PRECACHE & MODEL LOADERS (RUN BEFORE FIRST INFERENCE)
 # ===================================================================================
 
 def precache_assets() -> None:
@@ -262,9 +250,8 @@ def precache_assets() -> None:
     except Exception as e:
         print(f"Warning: GGUF pre-cache error: {e}")
 
-
-def _load_xtts(device: str = "cpu") -> Xtts:
-    """Load XTTS from the local cache. Keep CPU at startup to avoid ZeroGPU device mixups."""
+def _load_xtts(device: str) -> Xtts:
+    """Load XTTS from the local cache. Always CPU at startup for ZeroGPU compatibility."""
     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
@@ -283,20 +270,16 @@ def _load_xtts(device: str = "cpu") -> Xtts:
     print("XTTS model loaded.")
     return model
 
-
-def _load_llama() -> Llama:
-    """
-    Load Llama (Zephyr GGUF).
-    Keep simple & robust: default to CPU (works everywhere).
-    """
-    print("Loading LLM (Zephyr GGUF)...")
+def _load_llama_cpu_only() -> Llama:
+    """Load Llama (Zephyr GGUF) on CPU only (ZeroGPU friendly)."""
+    print("Loading LLM (Zephyr GGUF) on CPU...")
     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,     # CPU-only for reliability across Spaces/ZeroGPU
+        n_gpu_layers=0,   # never touch CUDA at startup
         n_ctx=4096,
         n_batch=512,
         verbose=False
@@ -304,27 +287,26 @@ def _load_llama() -> Llama:
     print("LLM loaded (CPU).")
     return llm
 
-
 def init_models_and_latents() -> None:
     """
-    Preload models on CPU and compute voice latents on CPU.
-    This avoids ZeroGPU's "mixed device" errors from torchaudio-based resampling.
+    Preload TTS and LLM on CPU and compute voice latents on CPU.
+    This avoids any CUDA tensors outside the @spaces.GPU window.
     """
     global tts_model, llm_model, voice_latents
 
+    target_device = "cpu"  # FORCE CPU at startup for ZeroGPU compatibility
+
     if tts_model is None:
-        tts_model = _load_xtts(device="cpu")  # always CPU at startup
+        tts_model = _load_xtts(device=target_device)
+    else:
+        tts_model.to("cpu")
 
     if llm_model is None:
-        llm_model = _load_llama()
+        llm_model = _load_llama_cpu_only()
 
+    # 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"),
@@ -333,21 +315,10 @@ def init_models_and_latents() -> None:
                 ("Thera", "thera-1.wav"),
             ]:
                 path = os.path.join("voices", filename)
-                gpt_lat, spk_emb = tts_model.get_conditioning_latents(
+                voice_latents[role] = tts_model.get_conditioning_latents(
                     audio_path=path, gpt_cond_len=30, max_ref_length=60
                 )
-                # 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.")
-
+        print("Voice latents ready (CPU).")
 
 # Ensure we close Llama cleanly to avoid __del__ issues at interpreter shutdown
 def _close_llm():
@@ -359,7 +330,6 @@ def _close_llm():
         pass
 atexit.register(_close_llm)
 
-
 # ===================================================================================
 # 4) INFERENCE HELPERS
 # ===================================================================================
@@ -367,17 +337,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_prompt = format_prompt_zephyr(prompt, history, system_message_text)
+    formatted = format_prompt_zephyr(prompt, history, system_message_text)
     stream = llm_instance(
-        formatted_prompt,
+        formatted,
         temperature=0.7,
         max_tokens=512,
         top_p=0.95,
         stop=LLM_STOP_WORDS,
         stream=True
     )
-    for response in stream:
-        ch = response["choices"][0]["text"]
+    for resp in stream:
+        ch = resp["choices"][0]["text"]
         try:
             is_single_emoji = (len(ch) == 1 and emoji.is_emoji(ch))
         except Exception:
@@ -386,29 +356,31 @@ 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[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)
-
+                          latents: Tuple[torch.Tensor, torch.Tensor]) -> Generator[bytes, None, None]:
+    gpt_cond_latent, speaker_embedding = latents
     try:
         for chunk in tts_instance.inference_stream(
             text=text,
             language=language,
-            gpt_cond_latent=gpt_cond_latent_t,
-            speaker_embedding=speaker_embedding_t,
+            gpt_cond_latent=gpt_cond_latent,
+            speaker_embedding=speaker_embedding,
             temperature=0.85,
         ):
             if chunk is None:
                 continue
-            # 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)
+            f32 = chunk.detach().cpu().numpy().squeeze().astype(np.float32)
+            f32 = np.clip(f32, -1.0, 1.0)
             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:
@@ -418,32 +390,34 @@ def generate_audio_stream(tts_instance: Xtts, text: str, language: str,
             except Exception:
                 pass
 
-
 # ===================================================================================
-# 5) ZERO-GPU ENTRYPOINT (safe on native GPU as well)
+# 5) ZERO-GPU ENTRYPOINT (also works on native GPU)
 # ===================================================================================
 
-@spaces.GPU(duration=120)  # GPU ops must occur inside this function when on ZeroGPU
+@spaces.GPU(duration=120)  # ZeroGPU allocates a GPU only for this function call
 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 (loaded on CPU)
+    # Ensure models/latents exist (CPU)
     if tts_model is None or llm_model is None or not voice_latents:
         init_models_and_latents()
 
-    # During the GPU window, move XTTS to CUDA if available; otherwise stay on CPU
+    # If ZeroGPU granted CUDA for this call, move XTTS to CUDA; keep LLM 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 kept CPU for simplicity & reliability)
+    # 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])
@@ -460,10 +434,13 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
         if not any(c.isalnum() for c in sentence):
             continue
 
-        audio_chunks = generate_audio_stream(tts_model, sentence, lang, voice_latents[chatbot_role])
+        # 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)
         pcm_data = b"".join(chunk for chunk in audio_chunks if chunk)
 
-        # Optional noise reduction (best-effort)
+        # Optional noise reduction (best-effort, CPU)
         try:
             data_s16 = np.frombuffer(pcm_data, dtype=np.int16)
             if data_s16.size > 0:
@@ -478,7 +455,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})
 
-    # Leave model on CPU after the ZeroGPU window
+    # Return XTTS to CPU to release GPU instantly
     try:
         tts_model.to("cpu")
     except Exception:
@@ -486,7 +463,6 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
 
     return results
 
-
 # ===================================================================================
 # 6) STARTUP: PRECACHE & UI
 # ===================================================================================
@@ -501,21 +477,17 @@ 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. Uses GPU only within the generation call when available.",
-        flagging_mode="never",
-        # removed allow_flagging (deprecated) to silence warning
+        description="Enter a prompt to generate a short story with voice narration using on-demand GPU.",
+        flagging_mode="never",      # avoid deprecated allow_flagging path
+        analytics_enabled=False,
     )
 
 if __name__ == "__main__":
-    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("===== 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("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")),
-        ssr_mode=False,  # disable experimental SSR noise
-    )
+    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))