SoundGame/tools/stable-audio-3/optimized/tflite/scripts/install.py

#!/usr/bin/env python3
"""Internal install plumbing — do NOT invoke directly. Use ../install.sh from the project root.

Called by install.sh after it has set up uv + .venv. Responsibilities:
1. (Optionally) pip-install requirements.txt — skipped when INSTALL_SKIP_PIP=1
   is set (install.sh sets this since uv already handled deps).
2. Ask which DiT bundles to download, then fetch them from HuggingFace.
3. Print example commands the user can copy-paste.

Unlike the MLX release, this TFLite/LiteRT stack is portable CPU — it runs on
macOS/Linux, x86/ARM. No Apple-Silicon requirement.
"""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

# Windows consoles default to a legacy code page (cp1252/cp437) that can't encode
# the ✓/→ status glyphs. Force UTF-8 with replacement; no-op on macOS/Linux.
for _stream in (sys.stdout, sys.stderr):
    try:
        _stream.reconfigure(encoding="utf-8", errors="replace")
    except AttributeError:
        pass   # exotic stream without reconfigure() — leave as-is

# Enable ANSI/VT escape processing in legacy Windows consoles (cmd.exe);
# harmless elsewhere and on modern Windows Terminal.
if os.name == "nt":
    os.system("")

# This file lives in <project>/scripts/. SCRIPT_DIR points at the project
# root (where models/, requirements.txt, .venv/, sa3 wrapper live).
SCRIPT_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(Path(__file__).resolve().parent))  # `from weights import …`

MIN_PY = (3, 9)
SKIP_PIP = os.environ.get("INSTALL_SKIP_PIP") == "1"


def step(msg: str) -> None:
    print(f"\n\033[1;36m→ {msg}\033[0m")


def check_environment() -> None:
    """Bail out early with a clear message if the interpreter is too old."""
    py = sys.version_info
    if (py.major, py.minor) < MIN_PY:
        print(
            f"\n\033[1;31merror\033[0m: this script is running under Python "
            f"{py.major}.{py.minor}, but this stack requires Python "
            f"{MIN_PY[0]}.{MIN_PY[1]}+.\n\n"
            f"Re-run install.py with a Python {MIN_PY[0]}.{MIN_PY[1]}+ interpreter, e.g.:\n"
            f"    /path/to/python3.11 install.py\n\n"
            f"On macOS:  brew install python@3.11\n"
            f"On Debian/Ubuntu:  apt install python3.11\n"
            f"On Windows:  install from https://www.python.org/downloads/ "
            f"(check 'Add python.exe to PATH')\n",
            file=sys.stderr,
        )
        sys.exit(1)
    # TFLite / LiteRT (ai_edge_litert) is portable across macOS/Linux, x86/ARM —
    # no platform gate here (unlike the Metal-backed MLX release).


def pip_install_requirements() -> None:
    if SKIP_PIP:
        step("Dependencies already installed by install.sh (uv) — skipping pip step")
        return
    step("Installing Python dependencies")
    req = SCRIPT_DIR / "requirements.txt"

    # Prefer pip; fall back to uv if pip isn't available in this interpreter
    # (e.g. user is running install.py directly inside a uv-created venv
    # that wasn't seeded with pip).
    pip_available = subprocess.run(
        [sys.executable, "-m", "pip", "--version"],
        capture_output=True,
    ).returncode == 0

    if pip_available:
        cmd = [sys.executable, "-m", "pip", "install", "-r", str(req)]
    elif _have("uv"):
        cmd = ["uv", "pip", "install", "--python", sys.executable, "-r", str(req)]
    else:
        print(
            f"\n\033[1;31merror\033[0m: neither pip nor uv is available in this "
            f"interpreter ({sys.executable}).\n"
            f"Install pip with `python -m ensurepip --upgrade`, or use ./install.sh "
            f"which sets everything up via uv.",
            file=sys.stderr,
        )
        sys.exit(1)
    print(f"  $ {' '.join(cmd)}")
    subprocess.check_call(cmd)


def _have(binary: str) -> bool:
    """True if the given binary is on PATH."""
    from shutil import which
    return which(binary) is not None


def main() -> None:
    import argparse
    ap = argparse.ArgumentParser(
        description="SA3 TFLite install — non-interactive bundle downloader + post-install help.",
        epilog="Normally invoked via ../install.sh, not directly.",
    )
    ap.add_argument("--download", default="",
                    metavar="BUNDLES",
                    help="Comma-separated list of bundles to pre-download "
                         "(sm-music, sm-sfx, medium). Without this flag, "
                         "nothing is downloaded — sa3_tflite.py will fetch any "
                         "missing weights from HuggingFace on first use.")
    cli = ap.parse_args()

    check_environment()
    pip_install_requirements()

    # Import after pip install so a fresh checkout works.
    from weights import DIT_BUNDLES, SHARED, BUNDLE_SIZES, bundle_status, ensure_local

    step("Current weights status")
    for name in DIT_BUNDLES:
        present, total = bundle_status(name)
        mark = "✓" if present == total else " "
        print(f"  [{mark}] {name:9s}  {present}/{total} files present   ({BUNDLE_SIZES[name]})")

    chosen: list[str] = []
    if cli.download.strip():
        chosen = [b.strip() for b in cli.download.split(",") if b.strip()]
        unknown = [b for b in chosen if b not in DIT_BUNDLES]
        if unknown:
            print(f"\n\033[1;31merror\033[0m: unknown bundle(s): {', '.join(unknown)}. "
                  f"Choices: {', '.join(DIT_BUNDLES)}", file=sys.stderr)
            sys.exit(1)

    if chosen:
        step(f"Downloading {len(chosen)} bundle(s): {', '.join(chosen)}")
        seen: set[str] = set()
        for name in chosen:
            print(f"\n[{name}]")
            items = DIT_BUNDLES[name] + SHARED
            for rel, _hf in items:
                if rel in seen:
                    continue
                seen.add(rel)
                ensure_local(rel)
    else:
        step("No --download set — weights will auto-download on first ./sa3 use")
        print(f"  To pre-download instead, pass:  {Path(sys.executable).name} install.py --download sm-music")
        print(f"                              or:  ./install.sh --download sm-music,medium")

    from examples import print_example_commands, BOLD, GREEN, RESET
    print_example_commands(header=f"{BOLD}{GREEN}✓ Install complete.{RESET}  Try these commands:")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nAborted.")
        sys.exit(130)