First commit

.gitignore

@@ -0,0 +1 @@
+/.venv

Makefile

@@ -23,23 +23,23 @@ PORT   ?= 7860
 # Core runtime deps (CPU-safe). Torch comes via transitive deps where needed;
 # you may pin torch externally if required by your environment.
 REQS = \
-    "numpy<2" \
-    "gradio==4.27.0" \
-    "python-dotenv" \
-    "huggingface_hub" \
-    "ffmpeg-python" \
-    "nltk" \
-    "emoji" \
-    "langid" \
-    "noisereduce" \
-    "TTS" \
-    "llama-cpp-python>=0.2.90"
+	"numpy<2" \
+	"gradio==4.27.0" \
+	"python-dotenv" \
+	"huggingface_hub" \
+	"ffmpeg-python" \
+	"nltk" \
+	"emoji" \
+	"langid" \
+	"noisereduce" \
+	"TTS" \
+	"llama-cpp-python>=0.2.90"
 
 # Dev tools (optional)
 DEV_REQS = \
-    "ruff" \
-    "black" \
-    "pip-tools"
+	"ruff" \
+	"black" \
+	"pip-tools"
 
 # ================================================================
 # Meta
@@ -113,12 +113,12 @@ check-ffmpeg:
 # ================================================================
 # Pre-download model assets and compute voice latents (runs your app's functions)
 precache: install check-ffmpeg
-	$(PY) - <<'PY'
-from app import precache_assets, init_models_and_latents
-precache_assets()
-init_models_and_latents()
-print("Precache complete.")
-PY
+	$(PY) - <<- 'PY'
+	from app import precache_assets, init_models_and_latents
+	precache_assets()
+	init_models_and_latents()
+	print("Precache complete.")
+	PY
 
 run: install
 	@echo "Starting app on port $(PORT)…"

app.py

@@ -2,7 +2,6 @@
 # 1) SETUP & IMPORTS
 # ===================================================================================
 from __future__ import annotations
-
 import os
 import sys
 import base64
@@ -10,34 +9,33 @@ 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."
     )
 
-# --- 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()
 
@@ -49,7 +47,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 ---
@@ -64,7 +62,6 @@ import langid
 import emoji
 import noisereduce as nr
 
-
 # ===================================================================================
 # 2) GLOBALS & HELPERS
 # ===================================================================================
@@ -72,12 +69,10 @@ import noisereduce as nr
 # NLTK data
 nltk.download("punkt", quiet=True)
 
-# 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]] = {}
+voice_latents: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
 
 # Config
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -87,6 +82,9 @@ 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. "
@@ -100,25 +98,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
@@ -135,13 +115,13 @@ def pcm_to_wav(pcm_data: bytes, sample_rate: int = 24000, channels: int = 1, bit
 
 def split_sentences(text: str, max_len: int) -> List[str]:
     sentences = nltk.sent_tokenize(text)
-    chunks: List[str] = []
-    for sent in sentences:
-        if len(sent) > max_len:
-            chunks.extend(textwrap.wrap(sent, max_len, break_long_words=True))
+    out: List[str] = []
+    for s in sentences:
+        if len(s) > max_len:
+            out.extend(textwrap.wrap(s, max_len, break_long_words=True))
         else:
-            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>"
@@ -151,7 +131,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:
     """
@@ -168,23 +147,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 any 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
@@ -194,14 +167,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:
@@ -234,9 +205,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
@@ -255,20 +225,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
@@ -276,27 +242,27 @@ 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
 
+    # Always CPU here (ZeroGPU rule)
+    target_device = "cpu"
+
     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"),
@@ -305,21 +271,11 @@ 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(
+                # Returns torch tensors; keep them on CPU
+                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():
@@ -331,7 +287,6 @@ def _close_llm():
         pass
 atexit.register(_close_llm)
 
-
 # ===================================================================================
 # 4) INFERENCE HELPERS
 # ===================================================================================
@@ -339,17 +294,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:
@@ -358,29 +313,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:
@@ -390,32 +347,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])
@@ -432,10 +391,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:
@@ -450,7 +412,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:
@@ -458,7 +420,6 @@ def generate_story_and_speech(secret_token_input: str, input_text: str, chatbot_
 
     return results
 
-
 # ===================================================================================
 # 6) STARTUP: PRECACHE & UI
 # ===================================================================================
@@ -473,16 +434,16 @@ 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",
+        description="Enter a prompt to generate a short story with voice narration using on-demand GPU.",
         allow_flagging="never",
+        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()