Background thread inject routes into running Gradio server

app.py

@@ -13,7 +13,6 @@ from fastapi.responses import JSONResponse
 from einops import rearrange
 from huggingface_hub import login
 from stable_audio_3 import StableAudioModel
-from contextlib import asynccontextmanager
 import urllib.request
 import urllib.parse
 import json as _json
@@ -133,15 +132,15 @@ API_RESOURCES = {
     },
     "web_search": {
         "name": "Web search",
-        "description": "Search the web via DuckDuckGo. Supports general, news, and image search. No API key required.",
+        "description": "Search the web via DuckDuckGo.",
         "endpoint": "/respite/search",
         "method": "GET",
         "input": {
             "q": "string (search query)",
             "backend": "string: 'text' (default), 'news', or 'images'",
             "max_results": "integer (1-50, default 10)",
-            "extract_top": "integer (0-5, default 0) - fetch full page content for top N results",
-            "region": "string (default 'wt-wt') - e.g. 'us-en', 'gb-en', 'de-de'",
+            "extract_top": "integer (0-5, default 0)",
+            "region": "string (default 'wt-wt')",
         },
         "output": "JSON with title, url, content per result.",
     },
@@ -168,15 +167,12 @@ def _get_runtime_specs():
     storage = shutil.disk_usage(os.getcwd())
     hardware = os.environ.get("SPACE_HARDWARE", "zero-a10g")
     is_zerogpu = "zero" in hardware.lower()
-    allocated_cores = _parse_cpu_count()
-    allocated_ram = _parse_ram_bytes()
-    host_cores = os.cpu_count()
     return {
         "platform": platform.platform(),
         "python_version": platform.python_version(),
-        "cpu_cores": allocated_cores,
-        "host_cpu_cores": host_cores,
-        "ram_bytes": allocated_ram,
+        "cpu_cores": _parse_cpu_count(),
+        "host_cpu_cores": os.cpu_count(),
+        "ram_bytes": _parse_ram_bytes(),
         "storage_total_bytes": storage.total,
         "storage_used_bytes": storage.used,
         "storage_free_bytes": storage.free,
@@ -202,12 +198,7 @@ try:
         with DDGS() as ddgs:
             if backend == "news":
                 hits = list(ddgs.news(query, max_results=max_results, region=region))
-                results = []
-                for r in hits:
-                    entry = {"title": r.get("title", ""), "url": r.get("url", ""), "content": r.get("body", ""), "source": r.get("source", ""), "date": r.get("date", "")}
-                    if r.get("image"):
-                        entry["image"] = r["image"]
-                    results.append(entry)
+                results = [{"title": r.get("title", ""), "url": r.get("url", ""), "content": r.get("body", ""), "source": r.get("source", ""), "date": r.get("date", "")} for r in hits]
             elif backend == "images":
                 hits = list(ddgs.images(query, max_results=max_results, region=region))
                 results = [{"title": r.get("title", ""), "url": r.get("image", ""), "source": r.get("source", ""), "thumbnail": r.get("thumbnail", ""), "width": r.get("width"), "height": r.get("height")} for r in hits]
@@ -226,24 +217,23 @@ try:
                                 if r["url"] == url:
                                     r["extracted_content"] = page_content[:8000]
                                     break
-                    except Exception as e:
-                        _log(f"Warning: failed to extract {url}: {e}")
+                    except Exception:
+                        pass
         return {"query": query, "backend": backend, "results": results, "total_results": len(results)}
 
 except ImportError:
-    _log("Warning: duckduckgo_search not installed, search disabled")
-    def _search_web(query, backend="text", max_results=10, extract_top=0, region="wt-wt"):
-        raise RuntimeError("Search not available: duckduckgo_search not installed")
+    def _search_web(query, **kw):
+        raise RuntimeError("Search not available")
 
 
 # ---------------------------------------------------------------------------
-# CPU LLM benchmark (llama.cpp)
+# CPU LLM benchmark
 # ---------------------------------------------------------------------------
 
 _BENCHMARK_MODELS = [
-    {"name": "Gemma 4 E4B", "repo": "unsloth/gemma-4-E4B-it-GGUF", "file": "gemma-4-E4B-it-Q4_K_M.gguf", "size_gb": 2.5, "prompt": "Explain recursion in programming in 3 sentences."},
-    {"name": "Qwen3 8B", "repo": "Qwen/Qwen3-8B-GGUF", "file": "Qwen3-8B-Q4_K_M.gguf", "size_gb": 5.0, "prompt": "Explain recursion in programming in 3 sentences."},
-    {"name": "Qwen3.8 27B", "repo": "unsloth/Qwen3.8-27B-GGUF", "file": "Qwen3.8-27B-Q4_0.gguf", "size_gb": 16.0, "prompt": "Explain recursion in programming in 3 sentences."},
+    {"name": "Gemma 4 E4B", "repo": "unsloth/gemma-4-E4B-it-GGUF", "file": "gemma-4-E4B-it-Q4_K_M.gguf", "size_gb": 2.5},
+    {"name": "Qwen3 8B", "repo": "Qwen/Qwen3-8B-GGUF", "file": "Qwen3-8B-Q4_K_M.gguf", "size_gb": 5.0},
+    {"name": "Qwen3.8 27B", "repo": "unsloth/Qwen3.8-27B-GGUF", "file": "Qwen3.8-27B-Q4_0.gguf", "size_gb": 16.0},
 ]
 _BENCHMARK_CACHE = os.path.join(tempfile.gettempdir(), "llm_benchmark")
 _BENCHMARK_RESULTS = []
@@ -261,12 +251,11 @@ def _ensure_llama_cpp():
     if os.path.isfile(llama_cli):
         _LLAMA_CPP_BIN = llama_cli
         return llama_cli
-    _log("Benchmark: cloning llama.cpp...")
+    _log("Benchmark: building llama.cpp from source...")
     src_dir = os.path.join(_BENCHMARK_CACHE, "llama_src")
     if os.path.isdir(src_dir):
         shutil.rmtree(src_dir, ignore_errors=True)
     subprocess.check_call(["git", "clone", "--depth=1", "https://github.com/ggml-org/llama.cpp.git", src_dir], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-    _log("Benchmark: building llama.cpp (this takes 2-5 minutes)...")
     n_threads = _parse_cpu_count() or 8
     build_dir = os.path.join(_BENCHMARK_CACHE, "llama_build")
     os.makedirs(build_dir, exist_ok=True)
@@ -278,9 +267,8 @@ def _ensure_llama_cpp():
     if candidates:
         shutil.copy2(candidates[0], llama_cli)
         os.chmod(llama_cli, 0o755)
-        _log(f"Benchmark: llama.cpp built at {llama_cli}")
     else:
-        raise RuntimeError(f"llama-cli not found after build in {build_dir}")
+        raise RuntimeError("llama-cli not found after build")
     _LLAMA_CPP_BIN = llama_cli
     return llama_cli
 
@@ -294,53 +282,37 @@ def _run_benchmark(model_filter=None):
     os.makedirs(_BENCHMARK_CACHE, exist_ok=True)
     llama_cli = _ensure_llama_cpp()
     n_threads = _parse_cpu_count() or 8
-    results = []
     prompt = "Explain recursion in programming in 3 sentences."
-    models = _BENCHMARK_MODELS
-    if model_filter:
-        models = [m for m in _BENCHMARK_MODELS if m["name"] == model_filter]
-        if not models:
-            raise ValueError(f"Unknown model '{model_filter}'. Available: {[m['name'] for m in _BENCHMARK_MODELS]}")
+    models = [m for m in _BENCHMARK_MODELS if m["name"] == model_filter] if model_filter else _BENCHMARK_MODELS
+    if model_filter and not models:
+        raise ValueError(f"Unknown model. Available: {[m['name'] for m in _BENCHMARK_MODELS]}")
+    results = []
     for model_info in models:
         result = {"name": model_info["name"], "status": "error"}
         try:
-            _log(f"Benchmark: downloading {model_info['name']} (~{model_info['size_gb']} GB)...")
             model_path = hf_hub_download(repo_id=model_info["repo"], filename=model_info["file"], cache_dir=_BENCHMARK_CACHE)
-            _log(f"Benchmark: running {model_info['name']} with {n_threads} threads...")
             env = os.environ.copy()
-            bin_dir = os.path.dirname(llama_cli)
-            env["LD_LIBRARY_PATH"] = bin_dir + ":" + env.get("LD_LIBRARY_PATH", "")
+            env["LD_LIBRARY_PATH"] = os.path.dirname(llama_cli)
             n_predict = 128 if model_info["size_gb"] <= 6 else 256
             timeout_sec = 300 if model_info["size_gb"] <= 6 else 600
             cmd = [llama_cli, "-m", model_path, "-t", str(n_threads), "-c", "2048", "-n", str(n_predict), "--temp", "0.7", "--top-p", "0.9", "-p", prompt, "--no-display-prompt", "--log-disable"]
             t0 = time.time()
-            try:
-                proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec, env=env)
-            except subprocess.TimeoutExpired:
-                result["error"] = "Timeout"
-                results.append(result)
-                continue
+            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec, env=env)
             wall_time = time.time() - t0
-            output_text = proc.stdout.strip()
-            _log(f"Benchmark: {model_info['name']} exit code: {proc.returncode}")
             prompt_tps = gen_tps = 0.0
             prompt_tok = gen_tok = 0
             for line in proc.stderr.splitlines():
-                tps_match = re.search(r"(\d+\.?\d*) tokens per second", line)
-                tok_match = re.search(r"/ (\d+) tokens", line)
+                tps_m = re.search(r"(\d+\.?\d*) tokens per second", line)
+                tok_m = re.search(r"/ (\d+) tokens", line)
                 if "prompt eval" in line.lower():
-                    if tps_match: prompt_tps = float(tps_match.group(1))
-                    if tok_match: prompt_tok = int(tok_match.group(1))
+                    if tps_m: prompt_tps = float(tps_m.group(1))
+                    if tok_m: prompt_tok = int(tok_m.group(1))
                 elif "eval time" in line.lower() and "prompt" not in line.lower():
-                    if tps_match: gen_tps = float(tps_match.group(1))
-                    if tok_match: gen_tok = int(tok_match.group(1))
-            if gen_tok == 0 and output_text:
-                gen_tok = len(output_text.split())
-            result.update({"status": "success", "total_time_sec": round(wall_time, 2), "prompt_tokens": prompt_tok, "completion_tokens": gen_tok, "prompt_tokens_per_sec": round(prompt_tps, 1), "generation_tokens_per_sec": round(gen_tps, 1), "output": output_text[:500]})
-            _log(f"Benchmark: {model_info['name']} -> {gen_tps:.1f} gen tok/s, {prompt_tps:.1f} prompt tok/s")
+                    if tps_m: gen_tps = float(tps_m.group(1))
+                    if tok_m: gen_tok = int(tok_m.group(1))
+            result.update({"status": "success", "total_time_sec": round(wall_time, 2), "prompt_tokens": prompt_tok, "completion_tokens": gen_tok, "prompt_tokens_per_sec": round(prompt_tps, 1), "generation_tokens_per_sec": round(gen_tps, 1), "output": proc.stdout.strip()[:500]})
         except Exception as e:
             result["error"] = str(e)
-            _log(f"Benchmark ERROR: {model_info['name']}: {e}")
         results.append(result)
         import glob as _glob
         for f in _glob.glob(os.path.join(_BENCHMARK_CACHE, "**", "*.gguf"), recursive=True):
@@ -352,7 +324,7 @@ def _run_benchmark(model_filter=None):
 
 
 # ---------------------------------------------------------------------------
-# Model cache & audio generation
+# Audio generation
 # ---------------------------------------------------------------------------
 
 MODEL_CACHE = {}
@@ -360,22 +332,18 @@ 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}', 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
 
 
@@ -384,12 +352,11 @@ def generate_audio(prompt, duration, steps, cfg_scale, seed, model_name):
 # ---------------------------------------------------------------------------
 
 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.")
+    gr.Markdown("# Respite API")
     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)
+            prompt = gr.Textbox(label="Prompt", placeholder="Describe the audio...", 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")
@@ -401,100 +368,72 @@ with gr.Blocks(title="Respite API") as demo:
 
 
 # ---------------------------------------------------------------------------
-# Route handlers
-# ---------------------------------------------------------------------------
-
-
-def _resources_handler():
-    return JSONResponse({"resources": API_RESOURCES})
-
-
-def _specs_handler():
-    return JSONResponse({**API_SPECS, "runtime": _get_runtime_specs()})
-
-
-def _search_handler(q: str = "", backend: str = "text", max_results: int = 10, extract_top: int = 0, region: str = "wt-wt"):
-    if not q:
-        return JSONResponse({"error": "Missing query parameter 'q'"}, status_code=400)
-    if backend not in ("text", "news", "images"):
-        return JSONResponse({"error": f"Invalid backend '{backend}'. Use 'text', 'news', or 'images'."}, status_code=400)
-    try:
-        results = _search_web(q, backend=backend, max_results=max_results, extract_top=extract_top, region=region)
-        return JSONResponse(results)
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=502)
-
-
-def _benchmark_handler(force: bool = False, model: str = ""):
-    try:
-        global _BENCHMARK_RESULTS
-        if force:
-            _BENCHMARK_RESULTS = []
-        model_filter = model.strip() if model.strip() else None
-        results = _run_benchmark(model_filter=model_filter)
-        return JSONResponse({
-            "status": "complete",
-            "models": results,
-            "system": {"cpu_cores": _parse_cpu_count(), "ram_bytes": _parse_ram_bytes(), "threads": _parse_cpu_count() or 8},
-        })
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-
-def _benchmark_test_handler():
-    import subprocess as _sp
-    try:
-        llama_cli = _ensure_llama_cpp()
-        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]
-        from huggingface_hub import hf_hub_download
-        model_path = hf_hub_download(repo_id="unsloth/gemma-4-E4B-it-GGUF", filename="gemma-4-E4B-it-Q4_K_M.gguf", cache_dir=_BENCHMARK_CACHE)
-        model_ok = os.path.isfile(model_path)
-        model_size = os.path.getsize(model_path) if model_ok else 0
-        return JSONResponse({"binary_path": llama_cli, "binary_exists": os.path.isfile(llama_cli), "binary_ok": binary_ok, "binary_info": binary_info, "model_path": model_path, "model_ok": model_ok, "model_size_gb": round(model_size / 1024**3, 2)})
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-
-# ---------------------------------------------------------------------------
-# After launch, inject routes into the running server.
+# After Gradio is launched, inject routes into the running server.
+#
+# Strategy: Gradio blocks while the server runs, so we start a thread that
+# waits for the server to be ready, then adds our routes to it.
 # ---------------------------------------------------------------------------
 
-import atexit
-
-def _inject_routes_after_launch():
-    """Wait for Gradio server to start, then add our routes."""
-    for attempt in range(30):
+def _inject_routes():
+    """Wait for Gradio server to start, then add API routes."""
+    for attempt in range(60):
         time.sleep(2)
-        server = getattr(demo, 'server', None)
+        # demo.server is set after launch() starts uvicorn
+        server = getattr(demo, "server", None)
         if server is None:
             continue
-        fa_app = getattr(server, 'app', None)
-        if fa_app is None:
+        app = getattr(server, "app", None)
+        if app is None:
             continue
-        _log(f"Injecting routes into running server (attempt {attempt+1})...")
+        _log(f"Injecting routes (attempt {attempt+1})...")
         try:
-            fa_app.add_api_route("/respite/resources", _resources_handler, methods=["GET"])
-            fa_app.add_api_route("/respite/specs", _specs_handler, methods=["GET"])
-            fa_app.add_api_route("/respite/search", _search_handler, methods=["GET"])
-            fa_app.add_api_route("/respite/benchmark", _benchmark_handler, methods=["GET"])
-            fa_app.add_api_route("/respite/benchmark/test", _benchmark_test_handler, methods=["GET"])
-            _log(f"Routes injected into running server on attempt {attempt+1}")
+            # Add all API routes
+            from fastapi import Request as _Req
+
+            @app.get("/respite/resources")
+            async def _resources():
+                return JSONResponse({"resources": API_RESOURCES})
+
+            @app.get("/respite/specs")
+            async def _specs():
+                return JSONResponse({**API_SPECS, "runtime": _get_runtime_specs()})
+
+            @app.get("/respite/search")
+            async def _search(q: str = "", backend: str = "text", max_results: int = 10, extract_top: int = 0, region: str = "wt-wt"):
+                if not q:
+                    return JSONResponse({"error": "Missing 'q' parameter"}, status_code=400)
+                try:
+                    return JSONResponse(_search_web(q, backend=backend, max_results=max_results, extract_top=extract_top, region=region))
+                except Exception as e:
+                    return JSONResponse({"error": str(e)}, status_code=502)
+
+            @app.get("/respite/benchmark")
+            async def _benchmark(force: bool = False, model: str = ""):
+                try:
+                    global _BENCHMARK_RESULTS
+                    if force:
+                        _BENCHMARK_RESULTS = []
+                    results = _run_benchmark(model_filter=model.strip() or None)
+                    return JSONResponse({"status": "complete", "models": results, "system": {"cpu_cores": _parse_cpu_count(), "ram_bytes": _parse_ram_bytes(), "threads": _parse_cpu_count() or 8}})
+                except Exception as e:
+                    return JSONResponse({"error": str(e)}, status_code=500)
+
+            _log("Routes injected into running server!")
             return
         except Exception as e:
-            _log(f"Route injection attempt {attempt+1} failed: {e}")
-    _log("Failed to inject routes after 30 attempts")
+            _log(f"Route injection failed: {e}")
+    _log("WARNING: Could not inject routes after 120s")
 
 
-_thread = threading.Thread(target=_inject_routes_after_launch, daemon=True)
-_thread.start()
+# Start route injection in background
+_inject_thread = threading.Thread(target=_inject_routes, daemon=True)
+_inject_thread.start()
 
-# Disable the Gradio Spaces file watcher thread which crashes because
-# importing `spaces` in a new thread after CUDA init fails.
+# Disable file watcher
 import gradio.utils as _gr_utils
 if hasattr(_gr_utils, "watchfn_spaces"):
     _gr_utils.watchfn_spaces = lambda *a, **kw: None
 
+_log("Starting Gradio server...")
 demo.queue(max_size=4, default_concurrency_limit=1)
-demo.launch(ssr_mode=False)
+demo.launch()