app.py

import os
import platform
import sys
import shutil
import tempfile
import traceback
import torch
import torchaudio
import gradio as gr
import spaces
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from einops import rearrange
from huggingface_hub import login
from stable_audio_3 import StableAudioModel

for _stream in (sys.stdout, sys.stderr):
    if hasattr(_stream, "reconfigure"):
        _stream.reconfigure(encoding="utf-8", errors="backslashreplace")

hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)

try:
    @spaces.GPU(duration=1)
    def _gpu_startup_touch():
        return "ok"
    _gpu_startup_touch()
except Exception as 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",
        "description": "Generate music or sound effects from a text prompt.",
        "endpoint": "/api/audio/generate",
        "method": "POST",
        "input": {
            "prompt": "string",
            "duration": "number (1-120 seconds)",
            "steps": "integer (1-50)",
            "cfg_scale": "number (0-10)",
            "seed": "integer (-1 for random)",
            "model": "small-music | small-sfx",
        },
        "output": "WAV audio file",
    }
}

API_SPECS = {
    "name": "Respite API",
    "version": "1.0.0",
    "description": "General-purpose AI API server with audio generation capabilities.",
    "base_path": "/respite",
    "authentication": "none",
    "content_types": ["application/json", "audio/wav"],
    "resources_endpoint": "/respite/resources",
    "specs_endpoint": "/respite/specs",
    "limits": {
        "max_concurrent_requests": 1,
        "max_queue_size": 4,
        "audio_max_duration_seconds": 120,
    },
}


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 = {}


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 Gradio's internal FastAPI server.
#
# The HF Spaces Gradio runner calls demo.launch() directly, so our own
# FastAPI app is never served.  We monkey-patch Blocks.launch so that
# AFTER Gradio creates its server, we inject routes onto that server.
# ---------------------------------------------------------------------------

_original_blocks_launch = gr.Blocks.launch


def _patched_blocks_launch(self, *args, **kwargs):
    _log("Intercepted Blocks.launch() - calling original...")
    try:
        result = _original_blocks_launch(self, *args, **kwargs)
    except Exception:
        _log("ERROR in original launch:")
        _log(traceback.format_exc())
        raise

    if self is demo:
        try:
            fa_app = getattr(self, "app", None)
            _log(f"demo.app is: {fa_app}")
            if fa_app is not None:
                @fa_app.get("/respite/resources")
                def _resources():
                    return JSONResponse({"resources": API_RESOURCES})

                @fa_app.get("/respite/specs")
                def _specs():
                    return JSONResponse({**API_SPECS, "runtime": _get_runtime_specs()})

                _log("Respite API routes injected successfully")
            else:
                _log("Warning: demo.app is None after launch")
        except Exception:
            _log("ERROR injecting routes:")
            _log(traceback.format_exc())

    return result


gr.Blocks.launch = _patched_blocks_launch

demo.queue(max_size=4, default_concurrency_limit=1)