Add audio.cpp runtime helpers

audiocpp_runtime.py

@@ -0,0 +1,557 @@
+"""audio.cpp runtime helpers for the Hugging Face Space.
+
+Everything audio.cpp-specific lives here so it can be imported both by the
+Gradio app and by user code in the Python tab:
+
+    import audiocpp_runtime as ac
+    ac.list_packages(task="tts")
+    result = ac.run(family="pocket_tts", task="tts", request={"text": "Hello"})
+"""
+from __future__ import annotations
+
+import base64
+import ctypes
+import glob
+import json
+import os
+import shutil
+import socket
+import subprocess
+import sys
+import tarfile
+import threading
+import time
+import urllib.error
+import urllib.request
+import uuid
+from pathlib import Path
+from typing import Any
+
+APP_DIR = Path(__file__).resolve().parent
+AUDIOCPP_VERSION = os.environ.get("AUDIOCPP_VERSION", "v0.9.0")
+ASSET_NAME = "audio-{tag}-bin-ubuntu-x64-cuda12.8-colab.tar.gz"
+RELEASE_URL = "https://github.com/0xShug0/audio.cpp/releases/download/{tag}/{asset}"
+
+HOME = Path(os.environ.get("AUDIOCPP_HOME", Path.home() / ".audiocpp"))
+BIN_ROOT = HOME / "bin"
+WORK_ROOT = HOME / "work"
+
+
+def _pick_models_root() -> Path:
+    override = os.environ.get("AUDIOCPP_MODELS_DIR")
+    if override:
+        return Path(override)
+    data = Path("/data")
+    if data.is_dir() and os.access(data, os.W_OK):
+        return data / "audiocpp-models"  # persistent storage, if the Space has it
+    return HOME / "models"
+
+
+MODELS_ROOT = _pick_models_root()
+STUB_SRC = APP_DIR / "stub" / "cuda_stub.c"
+STUB_DIR = HOME / "stub"
+
+# Spec task names -> audio.cpp runtime task tokens.
+TASK_TOKEN = {
+    "tts": "tts", "clone": "clon", "design": "vdes", "asr": "asr", "align": "align",
+    "vad": "vad", "diar": "diar", "sep": "sep", "vc": "vc", "svc": "svc", "s2s": "s2s",
+    "music": "gen", "sfx": "gen", "edit": "gen", "midi": "midi", "turn": "turn",
+    "speaker": "spk",
+}
+TASK_LABELS = {
+    "tts": "Text to speech",
+    "clone": "Voice cloning",
+    "design": "Voice design",
+    "asr": "Speech recognition (ASR)",
+    "align": "Forced alignment",
+    "vad": "Voice activity detection",
+    "diar": "Speaker diarization",
+    "sep": "Source separation",
+    "vc": "Voice conversion",
+    "svc": "Singing voice conversion",
+    "s2s": "Speech to speech / restoration",
+    "music": "Music generation",
+    "sfx": "Sound effects",
+    "edit": "Audio editing",
+    "midi": "Audio to MIDI / score",
+    "turn": "Turn detection",
+}
+# Families whose loader wants the package directory rather than one GGUF file.
+DIR_FAMILIES = {"auk", "lfm2_audio", "liveavatar", "minimax_music3", "yue2"}
+
+_setup_lock = threading.Lock()
+_install_lock = threading.Lock()
+_state: dict[str, Any] = {"ready": False, "error": None, "server": None}
+
+
+# --------------------------------------------------------------------------- setup
+def log(msg: str) -> None:
+    print(f"[audiocpp] {msg}", flush=True)
+
+
+def ensure_runtime() -> Path:
+    """Download and unpack the audio.cpp CUDA prebuilt once. Returns the bundle dir."""
+    with _setup_lock:
+        if _state["ready"]:
+            return _state["bundle"]
+        bundle = BIN_ROOT / AUDIOCPP_VERSION
+        server = bundle / "audiocpp_server"
+        if not server.exists():
+            bundle.mkdir(parents=True, exist_ok=True)
+            asset = ASSET_NAME.format(tag=AUDIOCPP_VERSION)
+            url = RELEASE_URL.format(tag=AUDIOCPP_VERSION, asset=asset)
+            archive = BIN_ROOT / asset
+            log(f"downloading {url}")
+            req = urllib.request.Request(url, headers={"User-Agent": "audiocpp-hf-space"})
+            with urllib.request.urlopen(req, timeout=600) as resp, open(archive, "wb") as out:
+                shutil.copyfileobj(resp, out, length=1 << 20)
+            with tarfile.open(archive, "r:gz") as tar:
+                tar.extractall(bundle, filter="data")
+            archive.unlink(missing_ok=True)
+            found = list(bundle.rglob("audiocpp_server"))
+            if not found:
+                raise RuntimeError("release archive has no audiocpp_server")
+            server = found[0]
+        server.chmod(0o755)
+        MODELS_ROOT.mkdir(parents=True, exist_ok=True)
+        WORK_ROOT.mkdir(parents=True, exist_ok=True)
+        build_cuda_stub()
+        _state.update(ready=True, bundle=server.parent, server=server)
+        log(f"runtime ready: {server}")
+        return server.parent
+
+
+def build_cuda_stub() -> None:
+    """Compile the tiny libcuda.so.1 stand-in used for CPU-only runs."""
+    target = STUB_DIR / "libcuda.so.1"
+    if target.exists():
+        return
+    STUB_DIR.mkdir(parents=True, exist_ok=True)
+    prebuilt = APP_DIR / "stub" / "libcuda.so.1"
+    if prebuilt.exists():
+        shutil.copy2(prebuilt, target)
+        return
+    cc = shutil.which("gcc") or shutil.which("cc")
+    if not cc:
+        log("no C compiler found; CPU fallback unavailable")
+        return
+    res = subprocess.run([cc, "-shared", "-fPIC", "-O2", "-Wl,-soname,libcuda.so.1",
+                          "-o", str(target), str(STUB_SRC)], capture_output=True, text=True)
+    if res.returncode != 0:
+        log("stub build failed: " + res.stderr[-500:])
+
+
+def cuda_lib_dirs() -> list[str]:
+    dirs: set[str] = set()
+    for base in sys.path:
+        for lib in glob.glob(os.path.join(base, "nvidia", "*", "lib")):
+            dirs.add(lib)
+    return sorted(dirs)
+
+
+def gpu_available() -> bool:
+    """True when a real NVIDIA driver and device are visible to this process."""
+    try:
+        ctypes.CDLL("libcuda.so.1")
+    except OSError:
+        return False
+    smi = shutil.which("nvidia-smi")
+    if smi:
+        try:
+            out = subprocess.run([smi, "-L"], capture_output=True, text=True, timeout=20)
+            return out.returncode == 0 and "GPU" in out.stdout
+        except Exception:
+            return False
+    return True
+
+
+def server_env(backend: str) -> dict[str, str]:
+    env = {k: v for k, v in os.environ.items()}
+    dirs = cuda_lib_dirs()
+    if backend == "cpu":
+        dirs.append(str(STUB_DIR))  # lets the CUDA build start without a driver
+    if env.get("LD_LIBRARY_PATH"):
+        dirs.append(env["LD_LIBRARY_PATH"])
+    env["LD_LIBRARY_PATH"] = ":".join(dirs)
+    env.setdefault("OMP_NUM_THREADS", str(os.cpu_count() or 4))
+    return env
+
+
+# --------------------------------------------------------------------------- catalog
+_catalog_cache: dict[str, dict] | None = None
+
+
+def specs_dir() -> Path:
+    return ensure_runtime() / "model_specs"
+
+
+def catalog() -> dict[str, dict]:
+    """family -> spec summary (tasks, packages, voices, options)."""
+    global _catalog_cache
+    if _catalog_cache is not None:
+        return _catalog_cache
+    out: dict[str, dict] = {}
+    for path in sorted(specs_dir().glob("*.json")):
+        try:
+            spec = json.loads(path.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+        family = spec.get("family") or path.stem
+        defaults = spec.get("package_defaults") or {}
+        pkgs = []
+        for p in spec.get("packages", []):
+            p = dict(p)
+            p.setdefault("download", defaults.get("download"))
+            pkgs.append(p)
+        ui = spec.get("ui") or {}
+        out[family] = {
+            "family": family,
+            "display_name": spec.get("display_name", family),
+            "description": spec.get("description", ""),
+            "category": spec.get("category", ""),
+            "status": spec.get("status", ""),
+            "tasks": spec.get("tasks", []),
+            "languages": spec.get("languages", []),
+            "packages": pkgs,
+            "recommended": ui.get("recommended_package"),
+            "voices": ui.get("builtin_voices") or [],
+            "default_voice": ui.get("default_voice"),
+            "request_options": (spec.get("options") or {}).get("request", []),
+            "docs": ui.get("docs", []),
+        }
+    _catalog_cache = out
+    return out
+
+
+def families_for_task(task: str) -> list[str]:
+    return sorted(f for f, s in catalog().items() if task in s["tasks"] and s["packages"])
+
+
+def package(pkg_id: str) -> tuple[dict, dict]:
+    for fam in catalog().values():
+        for p in fam["packages"]:
+            if p["id"] == pkg_id:
+                return fam, p
+    raise KeyError(f"unknown package {pkg_id!r}")
+
+
+def default_package(family: str) -> str | None:
+    fam = catalog()[family]
+    ids = [p["id"] for p in fam["packages"]]
+    if fam["recommended"] in ids:
+        return fam["recommended"]
+    for p in fam["packages"]:
+        if p.get("default"):
+            return p["id"]
+    gguf = [p["id"] for p in fam["packages"] if p.get("format") == "gguf"]
+    return (gguf or ids or [None])[0]
+
+
+def list_packages(task: str | None = None) -> list[dict]:
+    rows = []
+    for fam in catalog().values():
+        if task and task not in fam["tasks"]:
+            continue
+        for p in fam["packages"]:
+            rows.append({
+                "family": fam["family"], "package": p["id"], "format": p.get("format"),
+                "precision": p.get("precision"), "tasks": fam["tasks"],
+                "installed": is_installed(p["id"]),
+            })
+    return rows
+
+
+def package_dir(pkg: dict) -> Path:
+    return MODELS_ROOT / pkg["target_directory"]
+
+
+def model_path(pkg_id: str) -> Path:
+    """Path to hand to the loader, mirroring the audio.cpp WebUI's choice."""
+    fam, pkg = package(pkg_id)
+    root = package_dir(pkg)
+    if pkg.get("format") != "gguf" or fam["family"] in DIR_FAMILIES:
+        return root
+    files = pkg.get("files") or []
+    if fam["family"] == "minimax_h3":
+        want = "dit_int8.gguf" if "int8_dit" in pkg_id else "dit.gguf"
+        files = [f for f in files if f.lower().endswith("/" + want)]
+    ggufs = [f for f in files if f.lower().endswith(".gguf")]
+    if not ggufs:
+        return root
+    rel = ggufs[0].replace("\\", "/")
+    prefix = (pkg.get("strip_prefix") or "").rstrip("/")
+    if prefix and rel.startswith(prefix + "/"):
+        rel = rel[len(prefix) + 1:]
+    return root / rel
+
+
+def session_options_for(pkg_id: str) -> dict[str, str]:
+    fam, pkg = package(pkg_id)
+    if fam["family"] == "lfm2_audio":
+        backbone = next((f for f in pkg.get("files", []) if not Path(f).name.startswith("mmproj-")), None)
+        return {"lfm2_audio.model_gguf": Path(backbone).name} if backbone else {}
+    if fam["family"] == "minimax_music3":
+        table = {
+            "minimax_music3_q8_0": ("language_model_q8_0.gguf", "rvq_depth_decoder_q8_0.gguf", "transformer_q8_0.gguf"),
+            "minimax_music3_bf16": ("language_model_bf16.gguf", "rvq_depth_decoder_bf16.gguf", "transformer_bf16.gguf"),
+            "minimax_music3_q4_0": ("language_model_q4_0.gguf", "rvq_depth_decoder_q8_0.gguf", "transformer_q4_0.gguf"),
+        }
+        if pkg_id in table:
+            lm, rvq, flow = table[pkg_id]
+            return {"minimax_music3.language_model_gguf": lm,
+                    "minimax_music3.rvq_depth_decoder_gguf": rvq,
+                    "minimax_music3.flow_transformer_gguf": flow}
+    return {}
+
+
+def is_installed(pkg_id: str) -> bool:
+    try:
+        _, pkg = package(pkg_id)
+    except KeyError:
+        return False
+    root = package_dir(pkg)
+    if not root.exists():
+        return False
+    files = pkg.get("files") or []
+    prefix = (pkg.get("strip_prefix") or "").rstrip("/")
+    for f in files:
+        rel = f.replace("\\", "/")
+        if prefix and rel.startswith(prefix + "/"):
+            rel = rel[len(prefix) + 1:]
+        if not (root / rel).exists():
+            return False
+    return bool(files) or any(root.iterdir())
+
+
+def package_size(pkg_id: str) -> int | None:
+    """Download size in bytes from the Hugging Face API, if available."""
+    try:
+        from huggingface_hub import HfApi
+        _, pkg = package(pkg_id)
+        dl = pkg.get("download") or {}
+        if dl.get("kind") != "huggingface_snapshot" or not pkg.get("files"):
+            return None
+        infos = HfApi().get_paths_info(dl["repo"], pkg["files"], revision=dl.get("revision") or "main")
+        return sum(getattr(i, "size", 0) or 0 for i in infos)
+    except Exception:
+        return None
+
+
+def install(pkg_id: str, progress=None) -> Path:
+    """Download one package with the repo's model manager. Returns the model path."""
+    ensure_runtime()
+    if is_installed(pkg_id):
+        return model_path(pkg_id)
+    with _install_lock:
+        if is_installed(pkg_id):
+            return model_path(pkg_id)
+        bundle = ensure_runtime()
+        cmd = [sys.executable, str(bundle / "tools" / "model_manager_v2.py"),
+               "--specs-dir", str(bundle / "model_specs"), "install", pkg_id,
+               "--models-root", str(MODELS_ROOT), "--progress"]
+        log("installing " + pkg_id)
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        tail: list[str] = []
+        for line in proc.stdout:  # type: ignore[union-attr]
+            line = line.rstrip()
+            tail = (tail + [line])[-30:]
+            if progress and line.startswith("AUDIOCPP_PROGRESS"):
+                try:
+                    kv = dict(part.split("=", 1) for part in line.split()[1:] if "=" in part)
+                    done, total = int(kv.get("downloaded", 0)), int(kv.get("total", 0))
+                    if total > 0:
+                        progress(min(done / total, 1.0),
+                                 desc=f"Downloading {pkg_id} ({done / 1e9:.2f} / {total / 1e9:.2f} GB)")
+                except Exception:
+                    pass
+        if proc.wait() != 0:
+            raise RuntimeError(f"install of {pkg_id} failed:\n" + "\n".join(tail))
+    return model_path(pkg_id)
+
+
+def uninstall(pkg_id: str) -> None:
+    bundle = ensure_runtime()
+    subprocess.run([sys.executable, str(bundle / "tools" / "model_manager_v2.py"),
+                    "--specs-dir", str(bundle / "model_specs"), "uninstall", pkg_id,
+                    "--models-root", str(MODELS_ROOT)], check=False)
+
+
+def installed_packages() -> list[str]:
+    return [r["package"] for r in list_packages() if r["installed"]]
+
+
+# --------------------------------------------------------------------------- audio
+def to_wav(src: str | os.PathLike, dst_dir: str | os.PathLike, rate: int | None = None,
+           mono: bool = False, max_seconds: float | None = None) -> str:
+    """Convert any audio or video file to 16-bit PCM WAV with ffmpeg."""
+    dst = Path(dst_dir) / f"{Path(src).stem[:40]}-{uuid.uuid4().hex[:6]}.wav"
+    cmd = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", str(src), "-vn"]
+    if max_seconds:
+        cmd += ["-t", str(max_seconds)]
+    if rate:
+        cmd += ["-ar", str(rate)]
+    if mono:
+        cmd += ["-ac", "1"]
+    cmd += ["-c:a", "pcm_s16le", str(dst)]
+    res = subprocess.run(cmd, capture_output=True, text=True)
+    if res.returncode != 0:
+        raise RuntimeError(f"ffmpeg could not read {Path(src).name}: {res.stderr.strip()[-400:]}")
+    return str(dst)
+
+
+def input_format(task: str, family: str) -> tuple[int | None, bool]:
+    """Sample rate / mono policy matching the audio.cpp WebUI."""
+    if task == "sep" or family == "apollo":
+        return 44100, False
+    if task in {"asr", "vad", "diar", "align", "midi", "turn"}:
+        return 16000, True
+    return None, False
+
+
+# --------------------------------------------------------------------------- server
+def _free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _post(url: str, body: dict, timeout: float) -> dict:
+    data = json.dumps(body).encode()
+    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode("utf-8", "replace"))
+    except urllib.error.HTTPError as err:
+        payload = err.read().decode("utf-8", "replace")
+        try:
+            msg = json.loads(payload).get("error", {}).get("message", payload)
+        except Exception:
+            msg = payload
+        raise RuntimeError(f"HTTP {err.code}: {msg}") from None
+
+
+def _tail(path: Path, n: int = 60) -> str:
+    try:
+        return "\n".join(path.read_text(errors="replace").splitlines()[-n:])
+    except Exception:
+        return ""
+
+
+def _save_result(res: dict, out_dir: Path, stem: str) -> dict:
+    """Replace base64 payloads with files on disk."""
+    out: dict[str, Any] = {k: v for k, v in res.items()
+                           if k not in {"audio", "named_audio_outputs", "artifacts"}}
+    files: list[str] = []
+    if res.get("audio"):
+        p = out_dir / f"{stem}.wav"
+        p.write_bytes(base64.b64decode(res["audio"]))
+        files.append(str(p))
+        out["audio_file"] = str(p)
+    for item in res.get("named_audio_outputs") or []:
+        name = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(item.get("id", "out")))
+        p = out_dir / f"{stem}-{name}.wav"
+        p.write_bytes(base64.b64decode(item["audio"]))
+        files.append(str(p))
+        out.setdefault("named_audio_files", {})[item.get("id", name)] = str(p)
+    for art in res.get("artifacts") or []:
+        meta = art.get("meta") or {}
+        ext = meta.get("extension") or meta.get("ext") or {"midi": "mid", "video": "mp4", "json": "json",
+                                                             "text": "txt"}.get(art.get("kind", ""), "bin")
+        name = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(art.get("id", "artifact")))
+        p = out_dir / f"{stem}-{name}.{ext.lstrip('.')}"
+        p.write_bytes(base64.b64decode(art.get("payload", "")))
+        files.append(str(p))
+        out.setdefault("artifact_files", []).append({"id": art.get("id"), "kind": art.get("kind"),
+                                                     "meta": meta, "file": str(p)})
+    out["files"] = files
+    return out
+
+
+def execute(job: dict, prefer_gpu: bool = True) -> dict:
+    """Start audiocpp_server for one model, run every request in job, stop it.
+
+    job keys: family, task (spec name), model_path, requests (list of request
+    dicts), load_options, session_options, timeout (seconds), out_dir.
+    """
+    ensure_runtime()
+    started = time.time()
+    timeout = float(job.get("timeout") or 600)
+    backends = ["cuda", "cpu"] if (prefer_gpu and gpu_available()) else ["cpu"]
+    out_dir = Path(job.get("out_dir") or WORK_ROOT / uuid.uuid4().hex)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    errors: list[str] = []
+    for backend in backends:
+        port = _free_port()
+        model = {
+            "id": "m",
+            "family": job["family"],
+            "path": str(job["model_path"]),
+            "task": TASK_TOKEN.get(job["task"], job["task"]),
+            "mode": "offline",
+        }
+        if job.get("load_options"):
+            model["load_options"] = {k: str(v) for k, v in job["load_options"].items()}
+        if job.get("session_options"):
+            model["session_options"] = {k: str(v) for k, v in job["session_options"].items()}
+        cfg = {"host": "127.0.0.1", "port": port, "backend": backend, "device": 0,
+               "threads": os.cpu_count() or 4, "lazy_load": True,
+               "busy_timeout_ms": int(timeout * 1000), "models": [model]}
+        cfg_path = out_dir / f"server-{backend}.json"
+        cfg_path.write_text(json.dumps(cfg, indent=1))
+        log_path = out_dir / f"server-{backend}.log"
+        server = _state["server"]
+        with open(log_path, "w") as logf:
+            proc = subprocess.Popen([str(server), "--config", str(cfg_path), "--no-ui"],
+                                    cwd=str(server.parent), env=server_env(backend),
+                                    stdout=logf, stderr=subprocess.STDOUT)
+        try:
+            base = f"http://127.0.0.1:{port}"
+            deadline = time.time() + 60
+            while True:
+                if proc.poll() is not None:
+                    raise RuntimeError("server exited during startup")
+                try:
+                    with urllib.request.urlopen(base + "/health", timeout=2) as r:
+                        if r.status == 200:
+                            break
+                except Exception:
+                    pass
+                if time.time() > deadline:
+                    raise RuntimeError("server did not become ready")
+                time.sleep(0.3)
+            results = []
+            for i, request in enumerate(job["requests"]):
+                remaining = max(10.0, timeout - (time.time() - started))
+                t0 = time.time()
+                res = _post(base + "/v1/tasks/run", {"model": "m", "request": request}, remaining)
+                saved = _save_result(res, out_dir, request.get("id") or f"out{i + 1:02d}")
+                saved["wall_s"] = round(time.time() - t0, 2)
+                results.append(saved)
+            return {"backend": backend, "results": results, "log": _tail(log_path),
+                    "elapsed_s": round(time.time() - started, 2), "out_dir": str(out_dir)}
+        except Exception as err:
+            errors.append(f"[{backend}] {err}\n{_tail(log_path, 25)}")
+            # Request-level errors (HTTP 4xx from the model) would repeat on CPU,
+            # so only start-up / transport failures fall through to the CPU backend.
+            if str(err).startswith("HTTP 4"):
+                break
+        finally:
+            proc.terminate()
+            try:
+                proc.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+    raise RuntimeError("\n\n".join(errors))
+
+
+def run(family: str, task: str, request: dict | list[dict] | None = None, package_id: str | None = None,
+        load_options: dict | None = None, session_options: dict | None = None,
+        timeout: float = 600, prefer_gpu: bool = True, out_dir: str | None = None) -> dict:
+    """Convenience wrapper: install the package if needed and run request(s)."""
+    package_id = package_id or default_package(family)
+    path = install(package_id)
+    sess = {**session_options_for(package_id), **(session_options or {})}
+    reqs = request if isinstance(request, list) else [request or {}]
+    return execute({"family": family, "task": task, "model_path": str(path), "requests": reqs,
+                    "load_options": load_options or {}, "session_options": sess,
+                    "timeout": timeout, "out_dir": out_dir}, prefer_gpu=prefer_gpu)