app.py

"""audio.cpp on Hugging Face ZeroGPU.

A Gradio front end for https://github.com/0xShug0/audio.cpp: 100+ TTS, voice
cloning, ASR, diarization, separation, voice conversion and music models,
served by the native audio.cpp CUDA runtime. GPU work runs inside a ZeroGPU
call; when no GPU is available the same binary runs on CPU.
"""
from __future__ import annotations

import spaces  # import first so ZeroGPU can patch the process

import html
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import traceback
import urllib.parse
import urllib.request
import uuid
import zipfile
from pathlib import Path

import gradio as gr

import audiocpp_runtime as ac

APP_DIR = Path(__file__).resolve().parent
ALLOW_PYTHON = os.environ.get("ALLOW_PYTHON", "1").lower() not in {"0", "false", "no", "off"}
MAX_BATCH = int(os.environ.get("MAX_BATCH", "24"))
MAX_GPU_SECONDS = int(os.environ.get("MAX_GPU_SECONDS", "240"))
MAX_DOWNLOAD_BYTES = int(os.environ.get("MAX_DOWNLOAD_BYTES", str(300 * 1024 * 1024)))
SAMPLE_URL = "https://raw.githubusercontent.com/0xShug0/audio.cpp/main/assets/resources/{name}"
SAMPLE_DIR = ac.HOME / "samples"


def sample_file(name: str) -> str:
    """Fetch a small sample clip from the audio.cpp repo (cached)."""
    SAMPLE_DIR.mkdir(parents=True, exist_ok=True)
    dest = SAMPLE_DIR / name
    if not dest.exists():
        urllib.request.urlretrieve(SAMPLE_URL.format(name=name), dest)
    return str(dest)

# --------------------------------------------------------------------------- boot
SETUP_ERROR = ""
try:
    ac.ensure_runtime()
    CATALOG = ac.catalog()
except Exception as exc:  # keep the UI alive so the error is visible
    SETUP_ERROR = f"{type(exc).__name__}: {exc}"
    traceback.print_exc()
    CATALOG = {}

try:
    MODEL_PARAMS = json.loads((APP_DIR / "model_params.json").read_text(encoding="utf-8"))
except Exception:
    MODEL_PARAMS = {}

TASK_ORDER = ["tts", "clone", "design", "asr", "align", "diar", "vad", "sep", "vc", "svc",
              "s2s", "music", "sfx", "edit", "midi", "turn"]
TASKS = [t for t in TASK_ORDER if any(t in f["tasks"] for f in CATALOG.values())]
TASK_CHOICES = [(ac.TASK_LABELS.get(t, t), t) for t in TASKS]

TEXT_TASKS = {"tts", "clone", "design", "align", "music", "sfx", "edit", "s2s"}
AUDIO_TASKS = {"asr", "align", "vad", "diar", "sep", "vc", "svc", "s2s", "edit", "midi", "turn"}
VOICE_TASKS = {"tts", "clone", "design", "vc", "svc"}
INSTRUCT_TASKS = {"tts", "clone", "design", "s2s", "edit"}
MUSIC_TASKS = {"music", "sfx", "edit"}
TEXT_LABELS = {
    "tts": "Text to speak", "clone": "Text to speak", "design": "Text to speak",
    "align": "Transcript to align", "music": "Music prompt / caption", "sfx": "Sound prompt",
    "edit": "Edit instruction / prompt", "s2s": "Optional text prompt",
}

FAMILY_PREFS = {
    "tts": ["supertonic", "kokoro_tts", "pocket_tts", "omnivoice", "qwen3_tts"],
    "clone": ["pocket_tts", "omnivoice", "chatterbox", "qwen3_tts", "index_tts2"],
    "design": ["omnivoice", "qwen3_tts", "maya1"],
    "asr": ["qwen3_asr", "parakeet_tdt", "nemotron_asr", "sense_asr", "fun_asr_nano"],
    "align": ["qwen3_forced_aligner"],
    "diar": ["moss_transcribe_diarize", "nemotron_3_diar", "sortformer_diar"],
    "vad": ["silero_vad", "pulsevad"],
    "sep": ["htdemucs", "mel_band_roformer", "bs_roformer"],
    "vc": ["seed_vc", "chatterbox", "meanvc2"],
    "music": ["ace_step", "stable_audio", "heartmula"],
    "sfx": ["stable_audio", "midashenglm_gen"],
}


def families_for(task: str) -> list[tuple[str, str]]:
    fams = ac.families_for_task(task) if CATALOG else []
    prefs = [f for f in FAMILY_PREFS.get(task, []) if f in fams]
    ordered = prefs + [f for f in fams if f not in prefs]
    return [(f"{CATALOG[f]['display_name']} ({f})", f) for f in ordered]


def packages_for(family: str) -> list[tuple[str, str]]:
    if family not in CATALOG:
        return []
    rows = []
    for p in CATALOG[family]["packages"]:
        mark = " [downloaded]" if ac.is_installed(p["id"]) else ""
        fmt = p.get("format", "")
        label = f"{p.get('display_name') or p['id']} [{fmt}]{mark}"
        rows.append((label, p["id"]))
    return rows


def _human(n: int | None) -> str:
    if not n:
        return "unknown size"
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} PB"


def family_info(family: str, pkg: str | None) -> str:
    if family not in CATALOG:
        return ""
    f = CATALOG[family]
    langs = ", ".join(f["languages"][:30]) + (" ..." if len(f["languages"]) > 30 else "")
    docs = " · ".join(f"[{Path(d).name}](https://github.com/0xShug0/audio.cpp/blob/main/{d})" for d in f["docs"][:4])
    lines = [f"**{f['display_name']}** · tasks: `{', '.join(f['tasks'])}`" + (f" · status: {f['status']}" if f['status'] else ""),
             f["description"]]
    if langs:
        lines.append(f"Languages: {langs}")
    if f["voices"]:
        lines.append(f"Built-in voices: {len(f['voices'])}")
    if pkg:
        state = "downloaded" if ac.is_installed(pkg) else "downloads on first run"
        lines.append(f"Package `{pkg}`: {state}")
    if docs:
        lines.append(f"Docs: {docs}")
    return "\n\n".join(lines)


def param_defaults(family: str) -> tuple[str, str, str]:
    """(request options JSON, session options JSON, markdown reference table)."""
    req, sess = {}, {}
    for p in MODEL_PARAMS.get(family, []) if isinstance(MODEL_PARAMS.get(family), list) else []:
        default = p.get("default")
        if default in (None, ""):
            continue
        (sess if p.get("scope") == "session" else req)[p["name"]] = default
    rows = ["| option | type | default | description |", "|---|---|---|---|"]
    for o in CATALOG.get(family, {}).get("request_options", [])[:60]:
        desc = str(o.get("description", "")).replace("|", "/")
        rng = ""
        if "min" in o or "max" in o:
            rng = f" ({o.get('min', '')}..{o.get('max', '')})"
        vals = f" values: {', '.join(map(str, o['values']))}" if o.get("values") else ""
        rows.append(f"| `{o['name']}` | {o.get('type', '')}{rng} | {o.get('default', '')} | {desc}{vals} |")
    table = "\n".join(rows) if len(rows) > 2 else "_This model's spec declares no request options._"
    return (json.dumps(req, indent=2, ensure_ascii=False) if req else "{}",
            json.dumps(sess, indent=2, ensure_ascii=False) if sess else "{}",
            table)


# --------------------------------------------------------------------------- text input helpers
def extract_text(path: str) -> str:
    p = Path(path)
    ext = p.suffix.lower()
    if ext == ".pdf":
        from pypdf import PdfReader
        return "\n\n".join((page.extract_text() or "") for page in PdfReader(str(p)).pages)
    if ext == ".docx":
        import docx
        return "\n".join(par.text for par in docx.Document(str(p)).paragraphs)
    if ext == ".epub":
        parts = []
        with zipfile.ZipFile(p) as z:
            for name in sorted(z.namelist()):
                if name.lower().endswith((".xhtml", ".html", ".htm")):
                    parts.append(strip_html(z.read(name).decode("utf-8", "replace")))
        return "\n\n".join(parts)
    raw = p.read_bytes().decode("utf-8", "replace")
    if ext in {".html", ".htm", ".xhtml", ".xml"}:
        return strip_html(raw)
    if ext in {".srt", ".vtt"}:
        return strip_subtitles(raw)
    if ext == ".json":
        try:
            data = json.loads(raw)
            if isinstance(data, str):
                return data
            if isinstance(data, dict):
                for key in ("text", "input", "content", "transcript"):
                    if isinstance(data.get(key), str):
                        return data[key]
        except Exception:
            pass
    return raw


def strip_html(raw: str) -> str:
    raw = re.sub(r"(?is)<(script|style|head).*?</\1>", " ", raw)
    raw = re.sub(r"(?i)<br\s*/?>|</p>|</div>|</h\d>|</li>", "\n", raw)
    raw = re.sub(r"<[^>]+>", " ", raw)
    raw = html.unescape(raw)
    return re.sub(r"[ \t]+", " ", re.sub(r"\n\s*\n+", "\n\n", raw)).strip()


def strip_subtitles(raw: str) -> str:
    out = []
    for line in raw.splitlines():
        s = line.strip()
        if not s or s.isdigit() or "-->" in s or s.upper().startswith("WEBVTT"):
            continue
        out.append(re.sub(r"<[^>]+>", "", s))
    return " ".join(out)


def fetch_url(url: str, dest_dir: Path) -> Path:
    url = url.strip()
    if not re.match(r"^https?://", url):
        raise gr.Error("URLs must start with http:// or https://")
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (audio.cpp Space)"})
    with urllib.request.urlopen(req, timeout=60) as resp:
        ctype = resp.headers.get("Content-Type", "")
        name = Path(urllib.parse.urlparse(resp.geturl()).path).name or "download"
        if "." not in name:
            guess = {"text/html": ".html", "application/pdf": ".pdf", "audio/mpeg": ".mp3",
                     "audio/wav": ".wav", "text/plain": ".txt"}.get(ctype.split(";")[0].strip(), "")
            name += guess
        dest = dest_dir / f"{uuid.uuid4().hex[:6]}-{name[-80:]}"
        size = 0
        with open(dest, "wb") as out:
            while chunk := resp.read(1 << 20):
                size += len(chunk)
                if size > MAX_DOWNLOAD_BYTES:
                    raise gr.Error(f"Download is larger than {_human(MAX_DOWNLOAD_BYTES)}.")
                out.write(chunk)
    return dest


def load_text(files, url: str, current: str, mode: str) -> str:
    chunks = []
    tmp = Path(tempfile.mkdtemp(prefix="text-"))
    for f in files or []:
        path = f if isinstance(f, str) else getattr(f, "name", str(f))
        chunks.append(extract_text(path).strip())
    if url and url.strip():
        chunks.append(extract_text(str(fetch_url(url, tmp))).strip())
    if not chunks:
        raise gr.Error("Add a text file or a URL first.")
    loaded = "\n\n".join(c for c in chunks if c)
    if mode == "Append" and current and current.strip():
        return current.rstrip() + "\n\n" + loaded
    return loaded


def split_text(text: str, mode: str) -> list[str]:
    text = text.strip()
    if mode == "One clip per paragraph":
        parts = [p.strip() for p in re.split(r"\n\s*\n", text)]
    elif mode == "One clip per line":
        parts = [p.strip() for p in text.splitlines()]
    else:
        parts = [text]
    return [p for p in parts if p]


# --------------------------------------------------------------------------- GPU entry points
def _gpu_seconds(job: dict, seconds: int) -> int:
    return int(max(20, min(int(seconds or 90), MAX_GPU_SECONDS)))


@spaces.GPU(duration=_gpu_seconds)
def gpu_execute(job: dict, seconds: int) -> dict:
    return ac.execute(job, prefer_gpu=True)


def _python_seconds(code_path: str, seconds: int) -> int:
    return int(max(20, min(int(seconds or 90), MAX_GPU_SECONDS)))


@spaces.GPU(duration=_python_seconds)
def gpu_python(code_path: str, seconds: int) -> dict:
    return run_python_file(code_path, seconds)


def _looks_like_gpu_unavailable(err: Exception) -> bool:
    msg = str(err).lower()
    return any(k in msg for k in ("quota", "zerogpu", "no gpu", "gpu task aborted", "unlogged",
                                  "gpu is not available", "too many requests"))


# --------------------------------------------------------------------------- main run
def _parse_json(text: str, what: str) -> dict:
    text = (text or "").strip()
    if not text:
        return {}
    try:
        val = json.loads(text)
    except json.JSONDecodeError as err:
        raise gr.Error(f"{what} is not valid JSON: {err}")
    if not isinstance(val, dict):
        raise gr.Error(f"{what} must be a JSON object.")
    return val


def _collect_audio(mic, files, url, work: Path) -> list[str]:
    paths: list[str] = []
    if mic:
        paths.append(mic)
    for f in files or []:
        paths.append(f if isinstance(f, str) else getattr(f, "name", str(f)))
    if url and url.strip():
        for u in re.split(r"[\s,]+", url.strip()):
            if u:
                paths.append(str(fetch_url(u, work)))
    return paths


def _fmt_time(sample: int | float | None, rate: int | None) -> str:
    if sample is None or not rate:
        return "?"
    s = float(sample) / float(rate)
    return f"{int(s // 60):02d}:{s % 60:06.3f}"


def _srt_time(sec: float) -> str:
    ms = int(round(sec * 1000))
    h, ms = divmod(ms, 3600000)
    m, ms = divmod(ms, 60000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def describe_result(res: dict, label: str, out_dir: Path) -> tuple[str, str | None]:
    """Human-readable text for one result, plus an SRT file when timings exist."""
    lines = [f"### {label}"] if label else []
    rate = res.get("sample_rate")
    if res.get("text"):
        lines.append(res["text"])
    if res.get("language"):
        lines.append(f"_language: {res['language']}_")
    cues = []
    for key, title in (("speaker_turns", "Speaker turns"), ("segments", "Segments")):
        items = res.get(key) or []
        if items:
            lines.append(f"\n**{title}**")
            for it in items[:2000]:
                who = f" speaker {it['speaker_id']}:" if "speaker_id" in it else ""
                txt = f" {it['text']}" if it.get("text") else ""
                conf = f" ({it['confidence']:.2f})" if isinstance(it.get("confidence"), (int, float)) and not txt else ""
                lines.append(f"[{_fmt_time(it.get('start_sample'), rate)} - {_fmt_time(it.get('end_sample'), rate)}]{who}{txt}{conf}")
                if rate and it.get("text"):
                    cues.append((it["start_sample"] / rate, it["end_sample"] / rate,
                                 (f"[{it['speaker_id']}] " if "speaker_id" in it else "") + it["text"]))
    words = res.get("words") or []
    if words:
        shown = " ".join(f"{w.get('word')}@{_fmt_time(w.get('start_sample'), rate)}" for w in words[:400])
        lines.append(f"\n**Words** ({len(words)}): {shown}")
        if not cues and rate:
            group: list[dict] = []
            for w in words:
                group.append(w)
                if len(group) >= 10 or str(w.get("word", "")).endswith((".", "?", "!")):
                    cues.append((group[0]["start_sample"] / rate, group[-1]["end_sample"] / rate,
                                 " ".join(str(g.get("word", "")) for g in group)))
                    group = []
            if group:
                cues.append((group[0]["start_sample"] / rate, group[-1]["end_sample"] / rate,
                             " ".join(str(g.get("word", "")) for g in group)))
    if res.get("custom_schema_output"):
        lines.append("```json\n" + json.dumps(res["custom_schema_output"], indent=1, ensure_ascii=False)[:20000] + "\n```")
    srt = None
    if cues:
        srt_path = out_dir / f"{re.sub(r'[^A-Za-z0-9_-]+', '_', label or 'transcript')}.srt"
        srt_path.write_text("\n".join(f"{i}\n{_srt_time(a)} --> {_srt_time(b)}\n{t}\n"
                                      for i, (a, b, t) in enumerate(cues, 1)), encoding="utf-8")
        srt = str(srt_path)
    return "\n".join(lines), srt


def concat_wavs(paths: list[str], out: Path) -> str | None:
    if len(paths) < 2:
        return None
    lst = out.with_suffix(".txt")
    lst.write_text("".join(f"file '{p}'\n" for p in paths))
    res = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0",
                          "-i", str(lst), "-c:a", "pcm_s16le", str(out)], capture_output=True)
    return str(out) if res.returncode == 0 else None


def run_task(task, family, pkg, device, gpu_seconds, text, split_mode, audio_mic, audio_files, audio_url,
             voice_id, voice_ref, ref_text, instruct, lyrics, language, seed, duration, speaking_rate,
             extra_json, load_json, session_json, raw_json, progress=gr.Progress()):
    if SETUP_ERROR:
        raise gr.Error("audio.cpp failed to start: " + SETUP_ERROR)
    if not family or family not in CATALOG:
        raise gr.Error("Pick a model family.")
    pkg = pkg or ac.default_package(family)
    work = ac.WORK_ROOT / f"{time.strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:6]}"
    inputs = work / "inputs"
    inputs.mkdir(parents=True, exist_ok=True)
    status = []

    # ---- base request
    options = _parse_json(extra_json, "Model options")
    base: dict = {}
    if language and str(language).strip():
        base["language"] = str(language).strip()
    if task in VOICE_TASKS and voice_id and str(voice_id).strip():
        base["voice_id"] = str(voice_id).strip()
    if task in VOICE_TASKS and voice_ref:
        progress(0.02, desc="Preparing reference voice")
        base["voice_ref"] = ac.to_wav(voice_ref, inputs, mono=True, max_seconds=60)
    if ref_text and ref_text.strip() and task in VOICE_TASKS:
        base["reference_text"] = ref_text.strip()
    if instruct and instruct.strip() and task in INSTRUCT_TASKS:
        base["instruct"] = instruct.strip()
    if task in MUSIC_TASKS:
        if lyrics and lyrics.strip():
            base["lyrics"] = lyrics.strip()
        if duration and float(duration) > 0:
            base["duration_seconds"] = float(duration)
    if seed is not None and int(seed) >= 0:
        base["seed"] = int(seed)
    if speaking_rate and abs(float(speaking_rate) - 1.0) > 1e-6 and task in {"tts", "clone", "design"}:
        base["speaking_rate"] = float(speaking_rate)
    if options:
        base["options"] = options
    base.update(_parse_json(raw_json, "Raw request JSON"))

    if task == "clone" and "voice_ref" not in base:
        raise gr.Error("Voice cloning needs a reference voice recording.")
    if task in {"vc", "svc"} and "voice_ref" not in base:
        raise gr.Error("Voice conversion needs a target voice recording.")

    # ---- requests
    requests: list[dict] = []
    rate, mono = ac.input_format(task, family)
    audio_inputs = _collect_audio(audio_mic, audio_files, audio_url, inputs) if task in AUDIO_TASKS else []
    if task in AUDIO_TASKS and not audio_inputs and task not in {"s2s", "edit"} and "audio" not in base:
        raise gr.Error("Add input audio (upload, record, file, or URL).")
    texts = split_text(text or "", split_mode) if task in TEXT_TASKS else []
    if task in {"tts", "clone", "design", "music", "sfx", "align"} and not texts and "text" not in base:
        raise gr.Error("Enter some text (or load it from a file / URL).")

    if audio_inputs:
        for i, src in enumerate(audio_inputs[:MAX_BATCH]):
            progress(0.05 + 0.1 * i / max(1, len(audio_inputs)), desc=f"Converting {Path(src).name}")
            r = dict(base)
            r["audio"] = ac.to_wav(src, inputs, rate=rate, mono=mono)
            if texts:
                r["text"] = "\n\n".join(texts)
            r["id"] = re.sub(r"[^A-Za-z0-9_-]+", "_", Path(src).stem)[:40] + f"-{i + 1}"
            requests.append(r)
        if len(audio_inputs) > MAX_BATCH:
            status.append(f"Only the first {MAX_BATCH} audio inputs were processed.")
    elif texts:
        for i, t in enumerate(texts[:MAX_BATCH]):
            requests.append({**base, "text": t, "id": f"clip{i + 1:02d}"})
        if len(texts) > MAX_BATCH:
            status.append(f"Only the first {MAX_BATCH} text chunks were processed.")
    else:
        requests.append({**base, "id": "out01"})

    # ---- model download (CPU side, before the GPU is attached)
    if not ac.is_installed(pkg):
        progress(0.15, desc=f"Downloading {pkg}")
    model_path = ac.install(pkg, progress=progress)
    sess = {**ac.session_options_for(pkg), **_parse_json(session_json, "Session options")}
    job = {"family": family, "task": task, "model_path": str(model_path), "requests": requests,
           "load_options": _parse_json(load_json, "Load options"), "session_options": sess,
           "out_dir": str(work), "timeout": max(30, int(gpu_seconds) - 5)}

    # ---- run
    progress(0.6, desc="Running on ZeroGPU" if device == "gpu" else "Running on CPU")
    t0 = time.time()
    try:
        if device == "gpu":
            try:
                out = gpu_execute(job, int(gpu_seconds))
            except Exception as err:
                if not _looks_like_gpu_unavailable(err):
                    raise
                status.append(f"ZeroGPU unavailable ({str(err)[:160]}). Fell back to CPU.")
                job["timeout"] = 1800
                out = ac.execute(job, prefer_gpu=False)
        else:
            job["timeout"] = 1800
            out = ac.execute(job, prefer_gpu=False)
    except gr.Error:
        raise
    except Exception as err:
        raise gr.Error(str(err)[-3000:], duration=None)

    # ---- outputs
    files, audios, texts_out = [], [], []
    for req, res in zip(requests, out["results"]):
        files.extend(res.get("files", []))
        if res.get("audio_file"):
            audios.append(res["audio_file"])
        desc, srt = describe_result(res, req["id"] if len(requests) > 1 else "", work)
        if desc.strip():
            texts_out.append(desc)
        if srt:
            files.append(srt)
    combined = concat_wavs(audios, work / "combined.wav") if task in {"tts", "clone", "design"} else None
    if combined:
        files.insert(0, combined)
    main_audio = combined or (audios[0] if audios else next((f for f in files if f.endswith(".wav")), None))
    transcript = "\n\n".join(texts_out)
    if transcript:
        tpath = work / "result.md"
        tpath.write_text(transcript, encoding="utf-8")
        files.append(str(tpath))
    status.insert(0, f"Done on **{out['backend'].upper()}** in {time.time() - t0:.1f}s "
                     f"({len(requests)} request{'s' if len(requests) != 1 else ''}, package `{pkg}`).")
    details = {"backend": out["backend"], "elapsed_s": out["elapsed_s"], "package": pkg,
               "model_path": str(model_path), "requests": requests, "results": out["results"]}
    return (main_audio, files or None, transcript or "", details, "\n\n".join(status), out.get("log", ""),
            gr.update(choices=packages_for(family), value=pkg))


# --------------------------------------------------------------------------- UI callbacks
def on_task(task):
    fams = families_for(task)
    fam = fams[0][1] if fams else None
    is_text, is_audio = task in TEXT_TASKS, task in AUDIO_TASKS
    return (
        gr.update(choices=fams, value=fam),
        gr.update(visible=is_text, label=TEXT_LABELS.get(task, "Text")),
        gr.update(visible=is_text),
        gr.update(visible=is_audio),
        gr.update(visible=task in VOICE_TASKS),
        gr.update(visible=task in INSTRUCT_TASKS,
                  label="Voice description" if task == "design" else "Instruction / style (optional)"),
        gr.update(visible=task in MUSIC_TASKS),
        gr.update(visible=task in {"tts", "clone", "design"}),
    )


def on_family(family):
    if family not in CATALOG:
        return gr.update(choices=[], value=None), "", gr.update(), "{}", "{}", ""
    pkg = ac.default_package(family)
    f = CATALOG[family]
    voices = f["voices"]
    table = param_defaults(family)[2]
    return (gr.update(choices=packages_for(family), value=pkg),
            family_info(family, pkg),
            gr.update(choices=voices, value=f["default_voice"] if f["default_voice"] in voices else None),
            "{}", "{}", table)


def on_package(family, pkg):
    return family_info(family, pkg)


def check_size(pkg):
    if not pkg:
        return "Pick a package."
    size = ac.package_size(pkg)
    return f"`{pkg}`: {_human(size)} ({'downloaded' if ac.is_installed(pkg) else 'not downloaded yet'})"


# --------------------------------------------------------------------------- models tab
def catalog_rows():
    rows = []
    for f in sorted(CATALOG.values(), key=lambda x: x["family"]):
        inst = [p["id"] for p in f["packages"] if ac.is_installed(p["id"])]
        rows.append([f["family"], f["display_name"], ", ".join(f["tasks"]), len(f["packages"]),
                     ", ".join(inst), (f["description"] or "")[:160]])
    return rows


def all_package_choices():
    return [(f"{p['id']} ({f['family']})", p["id"]) for f in CATALOG.values() for p in f["packages"]]


def disk_report():
    used = 0
    for p in ac.MODELS_ROOT.rglob("*"):
        if p.is_file():
            used += p.stat().st_size
    free = shutil.disk_usage(ac.MODELS_ROOT).free if ac.MODELS_ROOT.exists() else 0
    return f"Models folder `{ac.MODELS_ROOT}` uses {_human(used)}, {_human(free)} free."


def models_download(pkg, progress=gr.Progress()):
    if not pkg:
        raise gr.Error("Pick a package.")
    ac.install(pkg, progress=progress)
    return f"Downloaded `{pkg}`.\n\n{disk_report()}", catalog_rows()


def models_delete(pkg):
    if not pkg:
        raise gr.Error("Pick a package.")
    ac.uninstall(pkg)
    return f"Removed `{pkg}`.\n\n{disk_report()}", catalog_rows()


# --------------------------------------------------------------------------- python tab
PY_EXAMPLE = '''# Helper module: audiocpp_runtime (imported as ac). Files you write to the
# current directory show up below (audio is previewed).
import audiocpp_runtime as ac

print("tasks:", sorted({t for r in ac.list_packages() for t in r["tasks"]}))
print("GPU visible:", ac.gpu_available())

result = ac.run(
    family="supertonic", task="tts",
    request={"text": "Hello from audio.cpp, running inside a Hugging Face Space."},
    out_dir=".",
)
print(result["backend"], result["results"][0]["files"])
'''

SECRET_HINTS = ("TOKEN", "SECRET", "PASSWORD", "API_KEY", "ACCESS_KEY", "PRIVATE")


def run_python_file(code_path: str, seconds: int) -> dict:
    workdir = Path(code_path).parent
    env = {k: v for k, v in os.environ.items() if not any(h in k.upper() for h in SECRET_HINTS)}
    env["PYTHONPATH"] = str(APP_DIR) + (":" + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
    env["MPLBACKEND"] = "Agg"
    t0 = time.time()
    try:
        proc = subprocess.run([sys.executable, "-u", code_path], cwd=workdir, env=env,
                              capture_output=True, text=True, timeout=max(10, int(seconds) - 3))
        out, code = proc.stdout + ("\n" + proc.stderr if proc.stderr else ""), proc.returncode
    except subprocess.TimeoutExpired as err:
        out = (err.stdout or "") + "\n" + (err.stderr or "") if isinstance(err.stdout, str) else ""
        out += f"\n[timed out after {seconds}s]"
        code = -1
    files = sorted(str(p) for p in workdir.rglob("*") if p.is_file() and p.name != "main.py"
                   and not p.name.startswith("server-"))
    return {"output": out[-50000:], "returncode": code, "files": files[:50], "seconds": round(time.time() - t0, 1)}


def run_python(code, on_gpu, seconds):
    if not ALLOW_PYTHON:
        raise gr.Error("The Python runner is disabled on this Space (ALLOW_PYTHON=0).")
    if not code or not code.strip():
        raise gr.Error("Write some code first.")
    workdir = ac.WORK_ROOT / f"py-{uuid.uuid4().hex[:8]}"
    workdir.mkdir(parents=True, exist_ok=True)
    code_path = workdir / "main.py"
    code_path.write_text(code, encoding="utf-8")
    if on_gpu:
        try:
            res = gpu_python(str(code_path), int(seconds))
        except Exception as err:
            if not _looks_like_gpu_unavailable(err):
                raise gr.Error(str(err)[-2000:])
            res = run_python_file(str(code_path), int(seconds))
            res["output"] = f"[ZeroGPU unavailable, ran on CPU: {str(err)[:160]}]\n" + res["output"]
    else:
        res = run_python_file(str(code_path), int(seconds))
    audio = next((f for f in res["files"] if f.lower().endswith((".wav", ".mp3", ".flac", ".ogg"))), None)
    header = f"exit code {res['returncode']} · {res['seconds']}s\n" + "-" * 40 + "\n"
    return header + res["output"], res["files"] or None, audio


# --------------------------------------------------------------------------- layout
CSS = """
#title {margin-bottom: 0}
.small-note {font-size: 0.85em; opacity: 0.8}
"""

TEXT_FILE_TYPES = [".txt", ".md", ".markdown", ".srt", ".vtt", ".json", ".csv", ".html", ".htm",
                   ".xml", ".pdf", ".docx", ".epub", ".rtf", ".log"]
AUDIO_FILE_TYPES = ["audio", "video", ".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".webm",
                    ".mp4", ".mkv", ".mov"]


def build_ui() -> gr.Blocks:
    first_task = TASKS[0] if TASKS else "tts"
    first_fams = families_for(first_task)
    first_fam = first_fams[0][1] if first_fams else None
    first_table = param_defaults(first_fam)[2] if first_fam else ""
    first_voices = CATALOG.get(first_fam, {}).get("voices", []) if first_fam else []

    with gr.Blocks(title="audio.cpp · ZeroGPU") as demo:
        gr.Markdown(
            "# audio.cpp on ZeroGPU\n"
            "Run [audio.cpp](https://github.com/0xShug0/audio.cpp) (a native C++/ggml engine for 100+ audio models) "
            "in the browser: TTS, voice cloning and design, ASR, alignment, diarization, VAD, stem separation, "
            "voice conversion, restoration, and music or sound generation. Models download on first use from "
            "[audio-cpp/audio.cpp-gguf](https://huggingface.co/audio-cpp/audio.cpp-gguf).",
            elem_id="title")
        if SETUP_ERROR:
            gr.Markdown(f"**Startup error:** `{SETUP_ERROR}`")

        with gr.Tabs():
            # ------------------------------------------------------------- run tab
            with gr.Tab("Run"):
                with gr.Row():
                    task = gr.Dropdown(TASK_CHOICES, value=first_task, label="Task", scale=2)
                    family = gr.Dropdown(first_fams, value=first_fam, label="Model family", scale=3)
                    pkg = gr.Dropdown(packages_for(first_fam) if first_fam else [],
                                      value=ac.default_package(first_fam) if first_fam else None,
                                      label="Package / precision", scale=3)
                with gr.Row():
                    device = gr.Radio([("ZeroGPU (CPU fallback)", "gpu"), ("CPU only", "cpu")], value="gpu",
                                      label="Device", scale=2)
                    gpu_seconds = gr.Slider(30, MAX_GPU_SECONDS, value=120, step=10,
                                            label="GPU time budget (seconds, includes model load)", scale=3)
                    size_btn = gr.Button("Check download size", scale=1)
                info = gr.Markdown(family_info(first_fam, ac.default_package(first_fam)) if first_fam else "")

                with gr.Row():
                    with gr.Column(scale=5):
                        text = gr.Textbox(label=TEXT_LABELS.get(first_task, "Text"), lines=7,
                                          placeholder="Type or paste text here, or load it from files / a URL below.",
                                          visible=first_task in TEXT_TASKS)
                        with gr.Accordion("Load text from files or a URL", open=False,
                                          visible=first_task in TEXT_TASKS) as text_loader:
                            text_files = gr.File(label="Text files (txt, md, srt, vtt, pdf, docx, epub, html, json)",
                                                 file_count="multiple", file_types=TEXT_FILE_TYPES)
                            text_url = gr.Textbox(label="Or a web page / text URL", placeholder="https://...")
                            with gr.Row():
                                load_mode = gr.Radio(["Replace", "Append"], value="Replace", label="Load mode")
                                load_btn = gr.Button("Load into text box")
                            split_mode = gr.Radio(["Whole text", "One clip per paragraph", "One clip per line"],
                                                  value="Whole text", label="Batch mode")

                        with gr.Group(visible=first_task in AUDIO_TASKS) as audio_group:
                            audio_mic = gr.Audio(sources=["upload", "microphone"], type="filepath",
                                                 label="Input audio (upload or record)")
                            audio_files = gr.File(label="More audio / video files (batch, any format ffmpeg reads)",
                                                  file_count="multiple", file_types=AUDIO_FILE_TYPES)
                            audio_url = gr.Textbox(label="Or audio / video URLs (space or comma separated)",
                                                   placeholder="https://.../speech.mp3")
                            sample_audio_btn = gr.Button("Use a sample speech clip", size="sm")

                        with gr.Group(visible=first_task in VOICE_TASKS) as voice_group:
                            voice_id = gr.Dropdown(first_voices, value=None, allow_custom_value=True,
                                                   label="Built-in voice id (optional)")
                            voice_ref = gr.Audio(sources=["upload", "microphone"], type="filepath",
                                                 label="Reference / target voice (cloning and conversion)")
                            ref_text = gr.Textbox(label="Reference transcript (helps some cloning models)", lines=2)
                            sample_voice_btn = gr.Button("Use a sample reference voice", size="sm")

                        instruct = gr.Textbox(label="Instruction / style (optional)", lines=2,
                                              visible=first_task in INSTRUCT_TASKS)
                        with gr.Group(visible=first_task in MUSIC_TASKS) as music_group:
                            lyrics = gr.Textbox(label="Lyrics (optional)", lines=5)
                            duration = gr.Number(label="Duration in seconds (0 = model default)", value=30)

                        with gr.Row():
                            language = gr.Textbox(label="Language (optional, e.g. en, zh, auto)", scale=2)
                            seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0, scale=1)
                            speaking_rate = gr.Slider(0.5, 2.0, value=1.0, step=0.05, label="Speaking rate",
                                                      visible=first_task in {"tts", "clone", "design"}, scale=2)

                        with gr.Accordion("Advanced model options", open=False):
                            gr.Markdown("Values here are sent as per-request options (overriding the fields above). "
                                        "The table lists every option the model's spec declares.",
                                        elem_classes="small-note")
                            extra_json = gr.Code("{}", language="json", label="Request options (JSON)", lines=6)
                            defaults_btn = gr.Button("Insert this model's suggested defaults", size="sm")
                            session_json = gr.Code("{}", language="json", label="Session options (JSON)", lines=3)
                            load_json = gr.Code("{}", language="json", label="Load options (JSON)", lines=3)
                            raw_json = gr.Code("{}", language="json",
                                               label="Raw request fields merged last (JSON, e.g. {\"emotion\": \"happy\"})",
                                               lines=3)
                            options_table = gr.Markdown(first_table)

                        run_btn = gr.Button("Run", variant="primary", size="lg")

                    with gr.Column(scale=4):
                        out_audio = gr.Audio(label="Output audio", type="filepath", interactive=False)
                        out_text = gr.Markdown(label="Text output")
                        out_files = gr.File(label="All output files (wav, stems, srt, midi, ...)", file_count="multiple",
                                            interactive=False)
                        out_status = gr.Markdown()
                        with gr.Accordion("Details (JSON)", open=False):
                            out_json = gr.JSON()
                        with gr.Accordion("Server log", open=False):
                            out_log = gr.Textbox(lines=12, max_lines=30, show_label=False)

                task.change(on_task, task, [family, text, text_loader, audio_group, voice_group, instruct,
                                            music_group, speaking_rate])
                family.change(on_family, family, [pkg, info, voice_id, extra_json, session_json, options_table])
                defaults_btn.click(lambda f: param_defaults(f)[:2] if f else ("{}", "{}"), family,
                                   [extra_json, session_json])
                pkg.change(on_package, [family, pkg], info)
                size_btn.click(check_size, pkg, info)
                sample_audio_btn.click(lambda: sample_file("sample_16k.wav"), None, audio_mic)
                sample_voice_btn.click(lambda: sample_file("b.wav"), None, voice_ref)
                load_btn.click(load_text, [text_files, text_url, text, load_mode], text)
                run_btn.click(
                    run_task,
                    [task, family, pkg, device, gpu_seconds, text, split_mode, audio_mic, audio_files, audio_url,
                     voice_id, voice_ref, ref_text, instruct, lyrics, language, seed, duration, speaking_rate,
                     extra_json, load_json, session_json, raw_json],
                    [out_audio, out_files, out_text, out_json, out_status, out_log, pkg],
                    concurrency_limit=2,
                )

            # ------------------------------------------------------------- models tab
            with gr.Tab("Models"):
                gr.Markdown("Browse every model family in the audio.cpp catalog and pre-download packages so the "
                            "first run starts faster. Downloads are lost when the Space restarts unless it has "
                            "persistent storage.")
                with gr.Row():
                    m_pkg = gr.Dropdown(all_package_choices(), label="Package", scale=4, filterable=True)
                    m_size = gr.Button("Size", scale=1)
                    m_dl = gr.Button("Download", variant="primary", scale=1)
                    m_rm = gr.Button("Delete", variant="stop", scale=1)
                m_status = gr.Markdown(disk_report() if CATALOG else "")
                m_table = gr.Dataframe(catalog_rows(), headers=["family", "name", "tasks", "packages", "downloaded",
                                                                "description"],
                                       interactive=False, wrap=True)
                m_size.click(check_size, m_pkg, m_status)
                m_dl.click(models_download, m_pkg, [m_status, m_table])
                m_rm.click(models_delete, m_pkg, [m_status, m_table])

            # ------------------------------------------------------------- python tab
            with gr.Tab("Python"):
                if ALLOW_PYTHON:
                    gr.Markdown("Run Python with the `audiocpp_runtime` helpers (`ac.run`, `ac.install`, "
                                "`ac.list_packages`, `ac.execute`, `ac.to_wav`). Code runs in a separate process with "
                                "a time limit and no Space secrets. Anything written to the working directory is "
                                "returned below.")
                    py_code = gr.Code(PY_EXAMPLE, language="python", lines=18, label="Code")
                    with gr.Row():
                        py_gpu = gr.Checkbox(value=True, label="Run inside a ZeroGPU call")
                        py_secs = gr.Slider(20, MAX_GPU_SECONDS, value=120, step=10, label="Time limit (seconds)")
                        py_btn = gr.Button("Run Python", variant="primary")
                    py_out = gr.Textbox(lines=16, max_lines=40, label="Output")
                    py_audio = gr.Audio(label="First audio file", type="filepath", interactive=False)
                    py_files = gr.File(label="Files", file_count="multiple", interactive=False)
                    py_btn.click(run_python, [py_code, py_gpu, py_secs], [py_out, py_files, py_audio],
                                 concurrency_limit=2)
                else:
                    gr.Markdown("The Python runner is disabled on this Space.")

            # ------------------------------------------------------------- about tab
            with gr.Tab("About"):
                gr.Markdown(f"""
**How it works.** The Space downloads the official audio.cpp `{ac.AUDIOCPP_VERSION}` CUDA 12.8 Linux build
(the same prebuilt used by the project's Colab notebook) plus NVIDIA's CUDA 12 runtime wheels. Each run
starts `audiocpp_server` with a one-model config inside a ZeroGPU call, sends the request to `/v1/tasks/run`,
then shuts it down. If ZeroGPU is unavailable (no quota, or the Space is on CPU hardware), the same binary
runs on CPU through a small `libcuda` stub, so everything still works, only slower.

**Tips**
- The GPU budget includes model loading. Small models (Supertonic, Kokoro, PocketTTS, Qwen3-ASR 0.6B,
  HTDemucs) finish in a few seconds; 4B to 8B models may need 120 s or more.
- Use **Batch mode** to turn each paragraph or line into its own clip; TTS clips are also joined into
  `combined.wav`. ASR, separation and the other audio tasks accept many files at once.
- Some families expose extra controls; open **Advanced model options** for the full option table.
- Packages marked `[safetensors]` may need a conversion step that this Space does not run. Prefer GGUF.

**Credits.** audio.cpp by [0xShug0](https://github.com/0xShug0) and contributors (Apache-2.0). Model weights
keep their original licenses; see the project's
[model license list](https://github.com/0xShug0/audio.cpp/blob/main/docs/model_licenses.md).
Per-model option defaults come from the project's WebUI config.
""")
    return demo


demo = build_ui()

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=4, max_size=64).launch(
        theme=gr.themes.Soft(), css=CSS, ssr_mode=False, show_error=True,
        allowed_paths=[str(ac.WORK_ROOT), str(SAMPLE_DIR)],
    )