"""Foundation-1 REST API wrapper server. Wraps RC-stable-audio-tools' generate_cond() behind FastAPI endpoints. Runs on port 8002. Managed as a subprocess by the KGOne gateway. Environment variables: FOUNDATION1_PRETRAINED_NAME HuggingFace model name (default: RoyalCities/Foundation-1) FOUNDATION1_CKPT_PATH Path to local .safetensors checkpoint (overrides pretrained name) FOUNDATION1_CONFIG_PATH Path to local model_config.json (required if CKPT_PATH is set) FOUNDATION1_SERVER_PORT Port to listen on (default: 8002) """ from __future__ import annotations import asyncio import json import logging import os import sys import time import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path from typing import Optional import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import FileResponse from pydantic import BaseModel # ── path setup ──────────────────────��──────────────────────────────��────────── # Must happen BEFORE importing stable_audio_tools — gradio.py opens config.json # relative to CWD at import time. ROOT_DIR = Path(__file__).parent.parent.resolve() FOUNDATION1_DIR = ROOT_DIR / "foundation1" OUTPUT_DIR = ROOT_DIR / "outputs" / "clip" if not FOUNDATION1_DIR.exists(): raise RuntimeError( f"Foundation-1 repo not found at {FOUNDATION1_DIR}. " "Run init.bat first to clone the submodules." ) os.chdir(FOUNDATION1_DIR) sys.path.insert(0, str(FOUNDATION1_DIR)) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) # ── imports from stable_audio_tools (CWD is now FOUNDATION1_DIR) ───────────── import stable_audio_tools.interface.gradio as _gradio_module # noqa: E402 from stable_audio_tools.interface.gradio import generate_cond, load_model # noqa: E402 # Override the output directory from the one baked into config.json _gradio_module.output_directory = str(OUTPUT_DIR) # ── logging ────────────────────────────────────────────────────────────────���── logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", ) logger = logging.getLogger(__name__) # ── configuration ───────────────────────────────────────────────────────────── PRETRAINED_NAME = os.environ.get("FOUNDATION1_PRETRAINED_NAME", "RoyalCities/Foundation-1") CKPT_PATH = os.environ.get("FOUNDATION1_CKPT_PATH") CONFIG_PATH = os.environ.get("FOUNDATION1_CONFIG_PATH") PORT = int(os.environ.get("FOUNDATION1_SERVER_PORT", "8002")) # Default local paths written by init.bat's snapshot_download step. # The HuggingFace repo names the file Foundation_1.safetensors (not model.safetensors), # so get_pretrained_model() would fail — we bypass it by using the local path directly. _DEFAULT_LOCAL_DIR = ROOT_DIR / "foundation1" / "models" / "RoyalCities-Foundation-1" _DEFAULT_CKPT = _DEFAULT_LOCAL_DIR / "Foundation_1.safetensors" _DEFAULT_CONFIG = _DEFAULT_LOCAL_DIR / "model_config.json" # ── in-memory task store ────────────────────────────────────────────────────── _tasks: dict[str, dict] = {} _executor = ThreadPoolExecutor(max_workers=1) # one generation at a time (single GPU) # ── app lifecycle ───────────────────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): logger.info("Loading Foundation-1 model...") _load_foundation1_model() logger.info("Foundation-1 model ready.") yield logger.info("Shutting down Foundation-1 server.") _gradio_module.model = None def _load_foundation1_model() -> None: import torch device = "cuda" if torch.cuda.is_available() else "cpu" logger.info("Using device: %s", device) # Resolve checkpoint: env var override > local download > pretrained_name fallback ckpt = Path(CKPT_PATH) if CKPT_PATH else (_DEFAULT_CKPT if _DEFAULT_CKPT.exists() else None) config = Path(CONFIG_PATH) if CONFIG_PATH else (_DEFAULT_CONFIG if _DEFAULT_CONFIG.exists() else None) if ckpt and config: with open(config) as f: model_config = json.load(f) load_model(model_config=model_config, model_ckpt_path=str(ckpt), device=device) logger.info("Loaded local checkpoint: %s", ckpt) else: # Last resort: let stable_audio_tools attempt its own download. # Note: get_pretrained_model() expects model.safetensors; this will fail unless # the upstream repo is updated to match that name. logger.warning( "Local weights not found at %s — falling back to pretrained_name download. " "Run init.bat to pre-download weights.", _DEFAULT_CKPT, ) load_model(pretrained_name=PRETRAINED_NAME, device=device) logger.info("Loaded pretrained model: %s", PRETRAINED_NAME) app = FastAPI( title="Foundation-1 Server", description="Wraps Foundation-1 (RC-stable-audio-tools) for REST API use.", version="0.1.0", lifespan=lifespan, ) # ── request model ───────────────────────────────────────────────────────────── class GenerateRequest(BaseModel): prompt: str negative_prompt: str = "" bars: int = 4 bpm: int = 140 note: str = "C" scale: str = "minor" steps: int = 75 cfg_scale: float = 7.0 seed: int = -1 sampler_type: str = "dpmpp-2m-sde" sigma_min: float = 0.03 sigma_max: float = 500.0 cfg_rescale: float = 0.0 # ── routes ──────────────────────────────────────────────────────────────────── @app.get("/health") async def health(): return {"status": "ok", "model_loaded": _gradio_module.model is not None} @app.post("/generate") async def generate(req: GenerateRequest): if _gradio_module.model is None: raise HTTPException(503, "Model not loaded") task_id = str(uuid.uuid4()) _tasks[task_id] = {"status": "pending", "created_at": time.time()} loop = asyncio.get_event_loop() loop.run_in_executor(_executor, _run_generate, task_id, req) return {"task_id": task_id} def _run_generate(task_id: str, req: GenerateRequest) -> None: _tasks[task_id]["status"] = "running" try: wav_path, _spectrograms, _piano_roll, midi_path = generate_cond( prompt=req.prompt, negative_prompt=req.negative_prompt or None, bars=req.bars, bpm=req.bpm, note=req.note, scale=req.scale, steps=req.steps, cfg_scale=req.cfg_scale, seed=req.seed, sampler_type=req.sampler_type, sigma_min=req.sigma_min, sigma_max=req.sigma_max, cfg_rescale=req.cfg_rescale, ) _tasks[task_id].update( { "status": "complete", "wav_filename": Path(wav_path).name if wav_path else None, "midi_filename": Path(midi_path).name if midi_path else None, } ) except Exception: logger.exception("Generation failed for task %s", task_id) _tasks[task_id]["status"] = "error" _tasks[task_id]["error"] = "Generation failed — check server logs." @app.get("/result/{task_id}") async def get_result(task_id: str): task = _tasks.get(task_id) if task is None: raise HTTPException(404, "Task not found") status = task["status"] if status != "complete": return {"task_id": task_id, "status": status, "error": task.get("error")} return {"task_id": task_id, "status": "complete"} @app.get("/audio/{task_id}") async def serve_audio(task_id: str): task = _tasks.get(task_id) if task is None: raise HTTPException(404, "Task not found") filename = task.get("wav_filename") if not filename: raise HTTPException(404, "No audio file for this task") path = OUTPUT_DIR / filename if not path.is_file(): raise HTTPException(404, "Audio file not found on disk") return FileResponse(str(path), media_type="audio/wav") @app.get("/midi/{task_id}") async def serve_midi(task_id: str): task = _tasks.get(task_id) if task is None: raise HTTPException(404, "Task not found") filename = task.get("midi_filename") if not filename: raise HTTPException(404, "No MIDI file for this task") path = OUTPUT_DIR / filename if not path.is_file(): raise HTTPException(404, "MIDI file not found on disk") return FileResponse(str(path), media_type="audio/midi") # ── entry point ────────────────────────────��────────────────────────────────���─ if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")