settled-reach/tooling/domains/assets/audio.py

"""Stable Audio Open connector — Gradio API wrapper.

Talks to the Stable Audio Open Gradio app on tower-of-joy :11500 (URL from
config.json). Uses the async Gradio pattern: POST to submit, then read the SSE
stream for the result, then download the file.

Formerly tooling/db/audio_connector.py, fronted by the audio-generate and
audio-health bash wrappers (T-1290). Behaviour is unchanged except where a
failure used to print `{"ok": false}` on stdout and exit 1: it now raises a
ReachError that says whether the service is OFF or rejected the request, and
post-processing is a function call rather than a second Python process.
"""

from __future__ import annotations

import json
import os
import shutil
import time
import urllib.error
import urllib.request

from tooling.core import console
from tooling.core.errors import ReachError
from tooling.domains.assets import audio_post, endpoints

SERVICE = "audio"


def get_base_url() -> str:
    return endpoints.base_url(SERVICE)


def health() -> dict:
    """Is the Stable Audio API reachable, and what does it expose?"""
    base = get_base_url()
    data = endpoints.call_json(
        urllib.request.Request(f"{base}/config", method="GET"),
        service=SERVICE,
        what="the health check",
        timeout=10,
    )
    api_names = [
        dep.get("api_name", "")
        for dep in data.get("dependencies", [])
        if dep.get("api_name", "") and not dep.get("api_name", "").startswith("js_")
    ]
    return {
        "ok": True,
        "url": base,
        "gradio_version": data.get("version", "unknown"),
        "api_endpoints": api_names,
    }


def post_process(wav_path: str, ogg_path: str | None = None, lufs=-16, quality=6, threshold=-50) -> dict:
    """trim + normalize + convert a generated WAV (the `--post` step)."""
    if ogg_path is None:
        ogg_path = os.path.splitext(wav_path)[0] + ".ogg"
    console.event(f"Post-processing → {os.path.basename(ogg_path)}...")
    try:
        return audio_post.pipeline(wav_path, ogg_path, quality=quality, lufs=lufs, threshold=threshold)
    except ReachError as exc:
        # Not fatal to the generation: the WAV exists and is reported, and the
        # failure is carried in the result exactly as before.
        return {"ok": False, "error": f"Post-processing failed: {exc.message}"}


def default_output(prompt: str, duration: float) -> str:
    """Auto-name: the first 40 prompt characters, sanitised, plus the duration."""
    safe = "".join(c if c.isalnum() or c in "-_ " else "" for c in prompt[:40])
    safe = safe.strip().replace(" ", "_").lower()
    return f"{safe}_{int(duration)}s.wav"


def extract_file_url(result_data, base: str) -> str | None:
    """Find the audio file URL in Gradio's `complete` payload, made absolute."""
    if isinstance(result_data, list) and len(result_data) > 0:
        audio_info = result_data[0]
    elif isinstance(result_data, dict) and "data" in result_data:
        audio_info = result_data["data"][0] if result_data["data"] else None
    else:
        audio_info = result_data

    file_url = None
    if isinstance(audio_info, dict):
        file_url = audio_info.get("url") or audio_info.get("path")
    elif isinstance(audio_info, str):
        file_url = audio_info
    if not file_url:
        return None

    if file_url.startswith("/"):
        return f"{base}{file_url}"
    if not file_url.startswith("http"):
        return f"{base}/file={file_url}"
    return file_url


def generate(
    prompt: str,
    duration: float = 10.0,
    steps: int = 100,
    cfg: float = 7.0,
    output: str | None = None,
    timeout: int = 600,
    post: bool = False,
    output_ogg: str | None = None,
) -> dict:
    """Generate audio from a text prompt; returns the result dict.

    duration is 0-47 s. Fewer steps is faster and worse. `post` (implied by
    `output_ogg`) runs trim + normalize + convert on the result.
    """
    base = get_base_url()
    # Fail fast, and say OFF rather than "submit failed", before building anything.
    endpoints.call(
        urllib.request.Request(f"{base}/config", method="GET"),
        service=SERVICE,
        what="the availability check",
        timeout=5,
    )
    api_url = f"{base}/gradio_api/call/generate_audio"
    output = output or default_output(prompt, duration)

    # 1. Submit.
    console.event(f"Submitting: {prompt!r} ({duration}s, {steps} steps, cfg {cfg})")
    submitted = endpoints.call_json(
        urllib.request.Request(
            api_url,
            data=json.dumps({"data": [prompt, duration, steps, cfg]}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        ),
        service=SERVICE,
        what="the generation request",
        timeout=30,
    )
    event_id = submitted.get("event_id") if isinstance(submitted, dict) else None
    if not event_id:
        raise ReachError(
            f"Stable Audio accepted the request but returned no event_id: {submitted}",
            fix="the Gradio app's API may have changed — check `reach assets audio health` endpoints",
        )
    console.event(f"Event {event_id} — waiting for generation (timeout {timeout}s)...")

    # 2. Read the SSE stream. A dropped connection is retried until the
    #    deadline, exactly as before; only an `error` event is fatal.
    started = time.time()
    result_data = _await_result(f"{api_url}/{event_id}", timeout, started)

    elapsed = round(time.time() - started, 1)
    console.event(f"Generation complete ({elapsed}s)")

    # 3. Download.
    file_url = extract_file_url(result_data, base)
    if not file_url:
        raise ReachError(
            f"could not find the audio URL in the response: {str(result_data)[:300]}",
            fix="the Gradio app's response shape may have changed — see audio.extract_file_url",
        )
    console.event(f"Downloading to {output}...")
    try:
        with urllib.request.urlopen(urllib.request.Request(file_url, method="GET"), timeout=60) as resp:
            with open(output, "wb") as f:
                shutil.copyfileobj(resp, f)
    except urllib.error.URLError as exc:
        raise ReachError(
            f"download failed from {file_url}: {exc}",
            fix="the file may have expired on the server — re-run the generation",
        ) from exc

    result = {
        "ok": True,
        "file": output,
        "size_bytes": os.path.getsize(output),
        "duration_requested": duration,
        "steps": steps,
        "cfg": cfg,
        "prompt": prompt,
        "generation_time_s": elapsed,
    }

    if post or output_ogg:
        post_result = post_process(output, ogg_path=output_ogg)
        if not post_result.get("ok"):
            result["post_processed"] = False
            result["post_error"] = post_result.get("error", "unknown")
        else:
            result["post_processed"] = True
            result["ogg_file"] = post_result.get("output", output_ogg)
            result["ogg_size_bytes"] = os.path.getsize(result["ogg_file"])

    return result


def _await_result(stream_url: str, timeout: int, start: float):
    """Poll the Gradio SSE stream until `complete`, an `error`, or the deadline."""
    last_status = None

    while time.time() - start < timeout:
        try:
            with urllib.request.urlopen(urllib.request.Request(stream_url, method="GET"), timeout=timeout) as resp:
                current_event = None
                for line_bytes in resp:
                    line = line_bytes.decode("utf-8").strip()
                    if line.startswith("event: "):
                        current_event = line[7:]
                        continue
                    if not (line.startswith("data: ") and current_event):
                        continue
                    data_str = line[6:]

                    if current_event == "heartbeat":
                        elapsed = int(time.time() - start)
                        if elapsed % 30 == 0 and elapsed > 0:
                            console.event(f"Still generating... ({elapsed}s elapsed)")
                    elif current_event == "error":
                        try:
                            detail = json.loads(data_str)
                        except (json.JSONDecodeError, TypeError):
                            detail = data_str
                        raise ReachError(
                            f"Stable Audio reported a generation error: {detail}",
                            fix="the service is up — check the prompt and duration (0-47 s), then re-run",
                        )
                    elif current_event == "complete":
                        try:
                            return json.loads(data_str)
                        except json.JSONDecodeError:
                            return data_str
                    elif current_event == "progress":
                        try:
                            status = str(json.loads(data_str))[:80]
                        except (json.JSONDecodeError, TypeError):
                            continue
                        if status != last_status:
                            console.event(f"Progress: {status}")
                            last_status = status
        except urllib.error.URLError:
            time.sleep(2)  # connection dropped mid-stream — retry until the deadline

    raise ReachError(
        f"generation timed out after {timeout}s",
        fix="raise --timeout, or lower --steps / --duration",
    )