Fix route injection: try demo.app directly, fallback to monkey-patch

app.py

@@ -54,7 +54,6 @@ def _read_cgroup_file(path):
 
 def _parse_cpu_count():
     """Read actual CPU allocation from cgroup limits."""
-    # cgroup v2: /sys/fs/cgroup/cpu.max -> "quota period" or "max period"
     cpu_max = _read_cgroup_file("/sys/fs/cgroup/cpu.max")
     if cpu_max:
         parts = cpu_max.split()
@@ -66,7 +65,6 @@ def _parse_cpu_count():
             except ValueError:
                 pass
 
-    # cgroup v1: /sys/fs/cgroup/cpu/cpu.cfs_quota_us and cpu.cfs_period_us
     quota = _read_cgroup_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
     period = _read_cgroup_file("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
     if quota and period:
@@ -77,7 +75,6 @@ def _parse_cpu_count():
         except ValueError:
             pass
 
-    # cpuset: /sys/fs/cgroup/cpuset.cpus -> "0-3" or "0,1,2,3"
     cpuset = _read_cgroup_file("/sys/fs/cgroup/cpuset.cpus")
     if cpuset:
         count = 0
@@ -95,28 +92,24 @@ def _parse_cpu_count():
 
 def _parse_ram_bytes():
     """Read actual memory limit from cgroup limits."""
-    # cgroup v2: /sys/fs/cgroup/memory.max
     mem_max = _read_cgroup_file("/sys/fs/cgroup/memory.max")
     if mem_max and mem_max != "max":
         try:
             val = int(mem_max)
-            # cgroup v2 reports in bytes
-            if val < 2**62:  # sanity check: not unlimited
+            if val < 2**62:
                 return val
         except ValueError:
             pass
 
-    # cgroup v1: /sys/fs/cgroup/memory/memory.limit_in_bytes
     mem_limit = _read_cgroup_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
     if mem_limit:
         try:
             val = int(mem_limit)
-            if val < 2**62:  # sanity check: not unlimited (reported as huge number)
+            if val < 2**62:
                 return val
         except ValueError:
             pass
 
-    # Fallback to /proc/meminfo
     try:
         with open("/proc/meminfo", "r", encoding="utf-8") as meminfo:
             for line in meminfo:
@@ -160,7 +153,7 @@ API_RESOURCES = {
             "region": "string (default 'wt-wt') - e.g. 'us-en', 'gb-en', 'de-de'",
         },
         "output": "JSON with title, url, content per result. News results include date and source. Extracted results include extracted_content.",
-    }
+    },
 }
 
 API_SPECS = {
@@ -185,7 +178,6 @@ def _get_runtime_specs():
     hardware = os.environ.get("SPACE_HARDWARE", "zero-a10g")
     is_zerogpu = "zero" in hardware.lower()
 
-    # Read actual container limits from cgroup, falling back to host values
     allocated_cores = _parse_cpu_count()
     allocated_ram = _parse_ram_bytes()
     host_cores = os.cpu_count()
@@ -216,15 +208,6 @@ try:
     from ddgs import DDGS
 
     def _search_web(query, backend="text", max_results=10, extract_top=0, region="wt-wt"):
-        """Search the web via DuckDuckGo.
-
-        Args:
-            query: Search query string.
-            backend: "text" (general), "news" (news articles), or "images".
-            max_results: Number of results to return (1-50).
-            extract_top: Fetch full page content for top N results (0-5).
-            region: DuckDuckGo region code (e.g. "us-en", "gb-en", "wt-wt" for worldwide).
-        """
         extract_top = max(0, min(5, extract_top))
         max_results = max(1, min(50, max_results))
 
@@ -267,7 +250,6 @@ try:
                         "content": r.get("body", ""),
                     })
 
-        # Optionally extract full page content from top results
         if extract_top > 0 and results and backend != "images":
             urls_to_extract = [r["url"] for r in results[:extract_top] if r.get("url")]
             with DDGS() as ddgs:
@@ -276,7 +258,6 @@ try:
                         extracted = ddgs.extract(url)
                         page_content = extracted.get("content", "")
                         if page_content:
-                            # Find the matching result and add extracted content
                             for r in results:
                                 if r["url"] == url:
                                     r["extracted_content"] = page_content[:8000]
@@ -299,94 +280,9 @@ except ImportError:
 
 
 # ---------------------------------------------------------------------------
-# Model cache
+# CPU LLM benchmark (llama.cpp)
 # ---------------------------------------------------------------------------
 
-MODEL_CACHE = {}
-
-
-def load_model(model_name):
-    if model_name not in MODEL_CACHE:
-        _log(f"Loading {model_name} model...")
-        model = StableAudioModel.from_pretrained(model_name, device="cpu")
-        MODEL_CACHE[model_name] = model
-        _log(f"{model_name} loaded successfully!")
-    return MODEL_CACHE[model_name]
-
-
-def generate_audio(prompt, duration, steps, cfg_scale, seed, model_name):
-    _log(
-        f"Generating with {model_name}: prompt='{prompt}', "
-        f"duration={duration}s, steps={steps}, cfg={cfg_scale}, seed={seed}"
-    )
-    model = load_model(model_name)
-    audio = model.generate(
-        prompt=prompt, duration=duration, steps=steps,
-        cfg_scale=cfg_scale, seed=seed, batch_size=1,
-    )
-    audio = rearrange(audio, "b d n -> d (b n)")
-    audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
-    output_path = os.path.join(
-        tempfile.gettempdir(),
-        f"stable_audio_{seed}_{hash(prompt) & 0xFFFFFFFF:08x}.wav",
-    )
-    torchaudio.save(output_path, audio, 44100)
-    _log("Generation complete!")
-    return output_path
-
-
-# ---------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------
-
-with gr.Blocks(title="Respite API") as demo:
-    gr.Markdown("# Respite API - Music & SFX Generation")
-    gr.Markdown(
-        "Generate music and sound effects using "
-        "Stability AI's Stable Audio 3 Small models."
-    )
-
-    with gr.Row():
-        with gr.Column():
-            model_name = gr.Dropdown(
-                choices=["small-music", "small-sfx"],
-                value="small-music", label="Model",
-            )
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Describe the music or sound effect you want to generate...",
-                lines=2,
-            )
-            duration = gr.Slider(minimum=1, maximum=120, value=30, step=1, label="Duration (seconds)")
-            steps = gr.Slider(minimum=1, maximum=50, value=8, step=1, label="Steps")
-            cfg_scale = gr.Slider(minimum=0.0, maximum=10.0, value=1.0, step=0.1, label="CFG Scale")
-            seed = gr.Number(value=-1, label="Seed (-1 for random)")
-            btn = gr.Button("Generate", variant="primary")
-
-        with gr.Column():
-            audio_output = gr.Audio(label="Generated Audio", type="filepath")
-
-    btn.click(
-        fn=generate_audio,
-        inputs=[prompt, duration, steps, cfg_scale, seed, model_name],
-        outputs=audio_output,
-    )
-
-
-# ---------------------------------------------------------------------------
-# Inject discovery routes into the Gradio FastAPI app.
-#
-# Gradio's Blocks.launch() creates a new FastAPI app via App.create_app()
-# and any routes on demo.app before launch() are lost.  We intercept the
-# app creation by wrapping the FastAPI app returned by the Gradio server's
-# configure_app / create_app lifecycle, then add our routes at startup.
-# ---------------------------------------------------------------------------
-
-from gradio.routes import App as GradioApp
-
-_orig_create_app = GradioApp.create_app
-
-# CPU LLM models for benchmarking via llama.cpp binary
 _BENCHMARK_MODELS = [
     {
         "name": "Gemma 4 E4B",
@@ -411,8 +307,8 @@ _BENCHMARK_MODELS = [
     },
 ]
 _BENCHMARK_CACHE = os.path.join(tempfile.gettempdir(), "llm_benchmark")
-_BENCHMARK_RESULTS = []  # cached after first run
-_LLAMA_CPP_BIN = None  # path to downloaded llama-cli binary
+_BENCHMARK_RESULTS = []
+_LLAMA_CPP_BIN = None
 
 
 def _ensure_llama_cpp():
@@ -431,12 +327,10 @@ def _ensure_llama_cpp():
         _LLAMA_CPP_BIN = llama_cli
         return llama_cli
 
-    # Clone and build llama.cpp from source
     _log("Benchmark: cloning llama.cpp...")
     src_dir = os.path.join(_BENCHMARK_CACHE, "llama_src")
     if os.path.isdir(src_dir):
-        import shutil as _shutil
-        _shutil.rmtree(src_dir, ignore_errors=True)
+        shutil.rmtree(src_dir, ignore_errors=True)
 
     subprocess.check_call(
         ["git", "clone", "--depth=1", "https://github.com/ggml-org/llama.cpp.git", src_dir],
@@ -456,12 +350,10 @@ def _ensure_llama_cpp():
         ["cmake", "--build", build_dir, "-j", str(n_threads), "--target", "llama-cli"],
         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
     )
-    # Find the built binary
     import glob as _glob
     candidates = _glob.glob(os.path.join(build_dir, "**", "llama-cli"), recursive=True)
     if candidates:
-        import shutil as _shutil2
-        _shutil2.copy2(candidates[0], llama_cli)
+        shutil.copy2(candidates[0], llama_cli)
         os.chmod(llama_cli, 0o755)
         _log(f"Benchmark: llama.cpp built at {llama_cli}")
     else:
@@ -472,16 +364,13 @@ def _ensure_llama_cpp():
 
 
 def _run_benchmark(model_filter=None):
-    """Run CPU benchmark on models using llama.cpp binary.
-
-    Args:
-        model_filter: If provided, only run this model name (e.g. 'Gemma 4 E4B').
-    """
+    """Run CPU benchmark on models using llama.cpp binary."""
     global _BENCHMARK_RESULTS
     if _BENCHMARK_RESULTS and not model_filter:
         return _BENCHMARK_RESULTS
 
     import subprocess
+    import re
     from huggingface_hub import hf_hub_download
 
     os.makedirs(_BENCHMARK_CACHE, exist_ok=True)
@@ -510,7 +399,7 @@ def _run_benchmark(model_filter=None):
             env = os.environ.copy()
             bin_dir = os.path.dirname(llama_cli)
             env["LD_LIBRARY_PATH"] = bin_dir + ":" + env.get("LD_LIBRARY_PATH", "")
-            # 27B model needs more tokens limit and longer timeout
+
             n_predict = 128 if model_info["size_gb"] <= 6 else 256
             timeout_sec = 300 if model_info["size_gb"] <= 6 else 600
             cmd = [
@@ -525,9 +414,7 @@ def _run_benchmark(model_filter=None):
                 "--no-display-prompt",
                 "--log-disable",
             ]
-            _log(f"Benchmark: cmd: {' '.join(cmd[:6])}...")
             _log(f"Benchmark: model exists: {os.path.isfile(model_path)}, size: {os.path.getsize(model_path) if os.path.isfile(model_path) else 0}")
-            _log(f"Benchmark: LD_LIBRARY_PATH={bin_dir}")
             t0 = time.time()
             try:
                 proc = subprocess.run(
@@ -541,20 +428,13 @@ def _run_benchmark(model_filter=None):
             output_text = proc.stdout.strip()
 
             _log(f"Benchmark: {model_info['name']} exit code: {proc.returncode}")
-            _log(f"Benchmark: {model_info['name']} stdout ({len(output_text)} chars): {output_text[:300]}")
-            _log(f"Benchmark: {model_info['name']} stderr ({len(proc.stderr)} chars): {proc.stderr[-600:]}")
+            _log(f"Benchmark: stderr tail: {proc.stderr[-600:]}")
 
-            # Parse llama.cpp stderr for timing info
-            # Various formats depending on version:
-            # "prompt eval time = 1234.56 ms / 12 tokens (102.88 ms per token, 9.72 tokens per second)"
-            # "eval time = 5678.90 ms / 64 tokens (88.73 ms per token, 11.27 tokens per second)"
             prompt_tps = 0.0
             gen_tps = 0.0
             prompt_tok = 0
             gen_tok = 0
-            import re
             for line in proc.stderr.splitlines():
-                # Match: "N.NN tokens per second"
                 tps_match = re.search(r"(\d+\.?\d*) tokens per second", line)
                 tok_match = re.search(r"/ (\d+) tokens", line)
                 if "prompt eval" in line.lower() or "prompt_eval" in line.lower():
@@ -568,10 +448,8 @@ def _run_benchmark(model_filter=None):
                     if tok_match:
                         gen_tok = int(tok_match.group(1))
 
-            # If parsing failed, try counting newlines as rough token estimate
             if gen_tok == 0 and output_text:
                 gen_tok = len(output_text.split())
-                _log(f"Benchmark: parsing failed, estimated {gen_tok} tokens from word count")
 
             result.update({
                 "status": "success",
@@ -590,7 +468,7 @@ def _run_benchmark(model_filter=None):
 
         results.append(result)
 
-        # Free disk cache between models to make room for next
+        # Free disk cache between models
         import glob as _glob
         for f in _glob.glob(os.path.join(_BENCHMARK_CACHE, "**", "*.gguf"), recursive=True):
             try:
@@ -603,8 +481,92 @@ def _run_benchmark(model_filter=None):
     return results
 
 
-def _patched_create_app(blocks, **kwargs):
-    fa_app = _orig_create_app(blocks, **kwargs)
+# ---------------------------------------------------------------------------
+# Model cache & audio generation
+# ---------------------------------------------------------------------------
+
+MODEL_CACHE = {}
+
+
+def load_model(model_name):
+    if model_name not in MODEL_CACHE:
+        _log(f"Loading {model_name} model...")
+        model = StableAudioModel.from_pretrained(model_name, device="cpu")
+        MODEL_CACHE[model_name] = model
+        _log(f"{model_name} loaded successfully!")
+    return MODEL_CACHE[model_name]
+
+
+def generate_audio(prompt, duration, steps, cfg_scale, seed, model_name):
+    _log(
+        f"Generating with {model_name}: prompt='{prompt}', "
+        f"duration={duration}s, steps={steps}, cfg={cfg_scale}, seed={seed}"
+    )
+    model = load_model(model_name)
+    audio = model.generate(
+        prompt=prompt, duration=duration, steps=steps,
+        cfg_scale=cfg_scale, seed=seed, batch_size=1,
+    )
+    audio = rearrange(audio, "b d n -> d (b n)")
+    audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    output_path = os.path.join(
+        tempfile.gettempdir(),
+        f"stable_audio_{seed}_{hash(prompt) & 0xFFFFFFFF:08x}.wav",
+    )
+    torchaudio.save(output_path, audio, 44100)
+    _log("Generation complete!")
+    return output_path
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+with gr.Blocks(title="Respite API") as demo:
+    gr.Markdown("# Respite API - Music & SFX Generation")
+    gr.Markdown(
+        "Generate music and sound effects using "
+        "Stability AI's Stable Audio 3 Small models."
+    )
+
+    with gr.Row():
+        with gr.Column():
+            model_name = gr.Dropdown(
+                choices=["small-music", "small-sfx"],
+                value="small-music", label="Model",
+            )
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe the music or sound effect you want to generate...",
+                lines=2,
+            )
+            duration = gr.Slider(minimum=1, maximum=120, value=30, step=1, label="Duration (seconds)")
+            steps = gr.Slider(minimum=1, maximum=50, value=8, step=1, label="Steps")
+            cfg_scale = gr.Slider(minimum=0.0, maximum=10.0, value=1.0, step=0.1, label="CFG Scale")
+            seed = gr.Number(value=-1, label="Seed (-1 for random)")
+            btn = gr.Button("Generate", variant="primary")
+
+        with gr.Column():
+            audio_output = gr.Audio(label="Generated Audio", type="filepath")
+
+    btn.click(
+        fn=generate_audio,
+        inputs=[prompt, duration, steps, cfg_scale, seed, model_name],
+        outputs=audio_output,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Add custom routes to demo.app AFTER the Blocks context closes.
+#
+# In Gradio 6.x, demo.app is set to a FastAPI app once the `with gr.Blocks()`
+# context exits.  We add our routes here.  When the HF runner calls
+# demo.launch(), it wraps demo.app — our routes survive because they're
+# already on the same object.
+# ---------------------------------------------------------------------------
+
+def _add_routes(fa_app):
+    """Add Respite API routes to a FastAPI app."""
 
     @fa_app.get("/respite/resources")
     def _resources():
@@ -628,7 +590,6 @@ def _patched_create_app(blocks, **kwargs):
 
     @fa_app.get("/respite/benchmark")
     def _benchmark(force: bool = False, model: str = ""):
-        """Run CPU LLM benchmark. Pass ?model=Gemma 4 E4B to run a single model."""
         try:
             global _BENCHMARK_RESULTS
             if force:
@@ -649,15 +610,12 @@ def _patched_create_app(blocks, **kwargs):
 
     @fa_app.get("/respite/benchmark/test")
     def _benchmark_test():
-        """Quick test: verify llama.cpp binary works."""
         import subprocess as _sp
         try:
             llama_cli = _ensure_llama_cpp()
-            # Test binary
             r = _sp.run([llama_cli, "--version"], capture_output=True, text=True, timeout=10)
             binary_ok = r.returncode == 0
             binary_info = r.stdout.strip()[:200] if binary_ok else r.stderr.strip()[:200]
-            # Test model download
             from huggingface_hub import hf_hub_download
             model_path = hf_hub_download(
                 repo_id="unsloth/gemma-4-E4B-it-GGUF",
@@ -678,11 +636,26 @@ def _patched_create_app(blocks, **kwargs):
         except Exception as e:
             return JSONResponse({"error": str(e)}, status_code=500)
 
-    _log("Respite API discovery routes injected")
-    return fa_app
+    _log("Respite API routes registered")
+
+
+# Try to add routes to demo.app (available after Blocks context exits)
+if hasattr(demo, "app") and demo.app is not None:
+    _add_routes(demo.app)
+    _log("Routes added to demo.app directly")
+else:
+    # demo.app might not exist yet. Use monkey-patch as fallback.
+    _log("demo.app not available, using monkey-patch fallback")
+
+    from gradio.routes import App as GradioApp
+    _orig_create_app = GradioApp.create_app
 
+    def _patched_create_app(blocks, **kwargs):
+        fa_app = _orig_create_app(blocks, **kwargs)
+        _add_routes(fa_app)
+        return fa_app
 
-GradioApp.create_app = staticmethod(_patched_create_app)
+    GradioApp.create_app = _patched_create_app
 
 # Disable the Gradio Spaces file watcher thread which crashes because
 # importing `spaces` in a new thread after CUDA init fails.
@@ -691,4 +664,4 @@ if hasattr(_gr_utils, "watchfn_spaces"):
     _gr_utils.watchfn_spaces = lambda *a, **kw: None
 
 demo.queue(max_size=4, default_concurrency_limit=1)
-demo.launch(ssr_mode=False)
+demo.launch()