Add API routes directly to Gradio FastAPI app

app.py

@@ -7,7 +7,6 @@ import torch
 import torchaudio
 import gradio as gr
 import spaces
-from fastapi import FastAPI
 
 # Hugging Face Spaces should use UTF-8; explicitly configure streams so model
 # libraries cannot inherit a platform-specific charmap encoding.
@@ -36,10 +35,28 @@ try:
         return "ok"
     _gpu_startup_touch()
 except Exception as e:
-    _log(f"Warning: ZeroGPU touch failed (running CPU-only): {e}")
+    print(f"Warning: ZeroGPU touch failed (running CPU-only): {e}", flush=True)
 
 
+def _log(message):
+    print(message.encode("ascii", "backslashreplace").decode("ascii"), flush=True)
+
+
+def _get_ram_bytes():
+    try:
+        with open("/proc/meminfo", "r", encoding="utf-8") as meminfo:
+            for line in meminfo:
+                if line.startswith("MemTotal:"):
+                    return int(line.split()[1]) * 1024
+    except (FileNotFoundError, OSError, ValueError):
+        pass
+    return None
+
+
+# ---------------------------------------------------------------------------
 # API metadata
+# ---------------------------------------------------------------------------
+
 API_RESOURCES = {
     "audio_generation": {
         "name": "Audio generation",
@@ -52,51 +69,12 @@ API_RESOURCES = {
             "steps": "integer (1-50)",
             "cfg_scale": "number (0-10)",
             "seed": "integer (-1 for random)",
-            "model": "small-music | small-sfx"
+            "model": "small-music | small-sfx",
         },
-        "output": "WAV audio file"
+        "output": "WAV audio file",
     }
 }
 
-def get_server_specs():
-    storage = shutil.disk_usage(os.getcwd())
-    return {
-        "name": "Respite API",
-        "version": "1.0.0",
-        "description": "General-purpose AI API server with audio generation capabilities.",
-        "base_path": "/api",
-        "authentication": "none",
-        "content_types": ["application/json", "audio/wav"],
-        "resources_endpoint": "/api/resources",
-        "specs_endpoint": "/api/specs",
-        "runtime": {
-            "platform": platform.platform(),
-            "python_version": platform.python_version(),
-            "cpu_cores": os.cpu_count(),
-            "ram_bytes": _get_ram_bytes(),
-            "storage_total_bytes": storage.total,
-            "storage_used_bytes": storage.used,
-            "storage_free_bytes": storage.free
-        },
-        "limits": {
-            "max_concurrent_requests": 1,
-            "max_queue_size": 4,
-            "audio_max_duration_seconds": 120
-        }
-    }
-
-
-def _get_ram_bytes():
-    try:
-        with open("/proc/meminfo", "r", encoding="utf-8") as meminfo:
-            for line in meminfo:
-                if line.startswith("MemTotal:"):
-                    return int(line.split()[1]) * 1024
-    except (FileNotFoundError, OSError, ValueError):
-        pass
-    return None
-
-
 API_SPECS = {
     "name": "Respite API",
     "version": "1.0.0",
@@ -109,24 +87,28 @@ API_SPECS = {
     "limits": {
         "max_concurrent_requests": 1,
         "max_queue_size": 4,
-        "audio_max_duration_seconds": 120
-    }
+        "audio_max_duration_seconds": 120,
+    },
 }
 
 
-def get_resources():
-    return {"resources": API_RESOURCES}
-
-
-def get_specs():
-    return {**API_SPECS, "runtime": get_server_specs()["runtime"]}
-
-
-def _log(message):
-    print(message.encode("ascii", "backslashreplace").decode("ascii"), flush=True)
+def _get_runtime_specs():
+    storage = shutil.disk_usage(os.getcwd())
+    return {
+        "platform": platform.platform(),
+        "python_version": platform.python_version(),
+        "cpu_cores": os.cpu_count(),
+        "ram_bytes": _get_ram_bytes(),
+        "storage_total_bytes": storage.total,
+        "storage_used_bytes": storage.used,
+        "storage_free_bytes": storage.free,
+    }
 
 
+# ---------------------------------------------------------------------------
 # Model cache
+# ---------------------------------------------------------------------------
+
 MODEL_CACHE = {}
 
 
@@ -134,10 +116,7 @@ def load_model(model_name):
     """Load model on demand and cache it."""
     if model_name not in MODEL_CACHE:
         _log(f"Loading {model_name} model...")
-        model = StableAudioModel.from_pretrained(
-            model_name,
-            device="cpu"
-        )
+        model = StableAudioModel.from_pretrained(model_name, device="cpu")
         MODEL_CACHE[model_name] = model
         _log(f"{model_name} loaded successfully!")
     return MODEL_CACHE[model_name]
@@ -148,7 +127,10 @@ def load_model(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}")
+    _log(
+        f"Generating with {model_name}: prompt='{prompt}', "
+        f"duration={duration}s, steps={steps}, cfg={cfg_scale}, seed={seed}"
+    )
 
     model = load_model(model_name)
 
@@ -158,82 +140,80 @@ def generate_audio(prompt, duration, steps, cfg_scale, seed, model_name):
         steps=steps,
         cfg_scale=cfg_scale,
         seed=seed,
-        batch_size=1
+        batch_size=1,
     )
 
     # Post-process: (batch, channels, samples) -> stereo waveform
     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")
+    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
 
 
-api = FastAPI(title=API_SPECS["name"], version=API_SPECS["version"])
-
-
-@api.get("/api/resources")
-def resources():
-    return JSONResponse(get_resources())
-
-
-@api.get("/api/specs")
-def specs():
-    return JSONResponse(get_specs())
-
+# ---------------------------------------------------------------------------
+# 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.")
+    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"
+                label="Model",
             )
             prompt = gr.Textbox(
                 label="Prompt",
                 placeholder="Describe the music or sound effect you want to generate...",
-                lines=2
+                lines=2,
             )
             duration = gr.Slider(
-                minimum=1, maximum=120, value=30, step=1,
-                label="Duration (seconds)"
+                minimum=1, maximum=120, value=30, step=1, label="Duration (seconds)"
             )
             steps = gr.Slider(
-                minimum=1, maximum=50, value=8, step=1,
-                label="Steps"
+                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)"
+                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"
-            )
+            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
+        outputs=audio_output,
     )
 
 
-# Mount the Gradio UI under the FastAPI application so API routes remain
-# available at /api/* instead of being handled by Gradio's catch-all route.
+# ---------------------------------------------------------------------------
+# Mount API discovery routes directly onto Gradio's FastAPI app.
+# This avoids conflicts with Gradio's internal /api proxy.
+# ---------------------------------------------------------------------------
+
 demo.queue(max_size=4, default_concurrency_limit=1)
-app = gr.mount_gradio_app(api, demo, path="/")
 
-# Spaces (sdk: python) serves the ASGI `app` variable directly.
-# For local development: uvicorn app:app --host 0.0.0.0 --port 7860
+
+@demo.app.get("/api/resources")
+def resources():
+    return JSONResponse({"resources": API_RESOURCES})
+
+
+@demo.app.get("/api/specs")
+def specs():
+    return JSONResponse({**API_SPECS, "runtime": _get_runtime_specs()})