Built a universal manager and installer for all operating systems and environments. Bumped a lot of versions of things. Still needs deep testing.

.gitignore

@@ -124,6 +124,9 @@ celerybeat.pid
 .venv
 .python
 .node
+.ffmpeg
+.mingit
+.uv
 env/
 venv/
 ENV/

README.md

@@ -81,7 +81,7 @@ cd ai-toolkit
 python3 -m venv venv
 source venv/bin/activate
 # install torch first
-pip3 install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
+pip3 install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
 pip3 install -r requirements.txt
 ```
 
@@ -97,7 +97,7 @@ git clone https://github.com/ostris/ai-toolkit.git
 cd ai-toolkit
 python -m venv venv
 .\venv\Scripts\activate
-pip install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
+pip install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
 pip install -r requirements.txt
 ```
 

dgx_instructions.md

@@ -37,7 +37,7 @@ conda activate ai-toolkit
 **2) Install PyTorch**
 
 ```
-pip3 install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu130
+pip3 install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
 ```
 
 

manager/README.md

@@ -0,0 +1,70 @@
+# AI Toolkit Manager
+
+Self-contained install/update manager for this checkout of AI Toolkit. Runs
+with any Python >= 3.8 and **no dependencies**, so it works before the
+training environment exists.
+
+```bash
+python3 -m manager install     # first-time setup: venv + torch + requirements
+python3 -m manager check       # is an update available / are deps out of sync?
+python3 -m manager update      # git pull, then sync deps + run migrations
+python3 -m manager launch      # start the web UI (http://localhost:8675)
+python3 -m manager doctor      # diagnose problems
+```
+
+## Design
+
+- **The install logic lives in the repo it installs.** Every commit knows how
+  to install itself; external frontends (the desktop launcher, `install.sh`,
+  `install.ps1`) just shell out to this CLI and stay dumb. Machine-readable
+  output via `check --json` / `detect --json`.
+- **Hardware → spec mapping** is in [spec.py](spec.py). One universal torch
+  pin (2.12.0 / torchvision 0.27.0 / torchaudio 2.11.0) on every platform:
+  cu130 wheels when the driver supports CUDA 13 (cu126 fallback for older
+  drivers, refused outright on Blackwell GPUs which need cu130), same stack +
+  Python 3.11 + `dgx_requirements.txt` on DGX/Grace, PyPI wheels on Mac,
+  rocm7.1 (experimental) for AMD, `--cpu` to force a CPU install. **Torch
+  pins there must be updated together with the README install instructions,
+  run_mac.zsh, and dgx_instructions.md.**
+- **Accelerators everywhere wheels exist**, via per-spec `extra_packages`
+  (installed after requirements with `--upgrade` so they override pins) and
+  `optional_packages` (installed one-by-one, warn-only on failure):
+  `torchcodec==0.15.0` on all platforms; flash-attn 2.8.3 prebuilt wheels
+  (mjun0812) on Linux x86_64/aarch64 + Windows; NATTEN 0.21.7 wheels
+  (whl.natten.org) on Linux both arches; triton bundled with torch on Linux
+  and `triton-windows` 3.7.x on Windows. No flash-attn/NATTEN/triton on Mac,
+  no NATTEN on Windows (no wheels exist).
+- **Nothing global is ever installed.** FFmpeg (shared builds — the libs
+  torchcodec dlopens) goes to `.ffmpeg/` ([ffmpeg.py](ffmpeg.py)), Node
+  (when the system lacks >= 20) to `.node/` ([nodejs.py](nodejs.py)), the uv
+  binary (when absent) to `.uv/` ([uvbin.py](uvbin.py)) with uv-managed
+  Pythons kept in `.uv/python/` via `UV_PYTHON_INSTALL_DIR`, and on Windows
+  without git, portable MinGit to `.mingit/` ([gitwin.py](gitwin.py)) — all
+  inside the repo and gitignored. The first clone on a git-less Windows
+  box is handled by the bootstrap layer (install.ps1 / desktop launcher),
+  which downloads MinGit itself and moves it into the checkout afterwards. `manager launch` puts them on PATH (and
+  LD_LIBRARY_PATH on Linux) for the whole UI/training process tree, and a
+  generated `sitecustomize.py` in the venv exposes ffmpeg to any direct use
+  of the venv python (plus `os.add_dll_directory` on Windows).
+- **Hostile-environment hardening** (learned from the community Windows
+  installer): every python/pip subprocess runs with PYTHONPATH/PYTHONHOME/
+  CONDA/PYENV/PIP_* scrubbed from the env; git runs with
+  `GIT_LFS_SKIP_SMUDGE=1`; git-pinned requirements (diffusers) are
+  force-reinstalled when requirements change since pip skips unchanged
+  version numbers; `launch` polls the UI port and opens the browser when
+  ready (`--no-browser` to disable, auto-skipped on headless boxes).
+- **uv is used when present** (fast installs, auto-downloads the right
+  Python); plain `venv` + `pip` otherwise. The venv is created at `.venv/`
+  (an existing `venv/` is also respected, matching `ui/cron/pythonPath.ts`).
+- **State** (requirements hash, applied migrations) lives inside the venv
+  (`aitk_manager_state.json`) — deleting the venv resets everything.
+- **Update flow**: `update` pulls fast-forward only, then **re-execs**
+  `python -m manager sync` so the freshly pulled manager code — not the stale
+  in-memory copy — performs its own dependency sync and migrations.
+  **Local work is never overwritten**: a dirty tree aborts the update by
+  default (untracked files don't count), `--auto` (used by the run_* scripts)
+  warns and skips the pull instead so launching still works, and there is no
+  reset/clean anywhere — even `--force` relies on git itself refusing to
+  clobber modified files.
+- **Migrations** ([migrations.py](migrations.py)): one-time post-update steps,
+  each applied at most once per environment.

manager/__init__.py

@@ -0,0 +1,6 @@
+"""AI Toolkit install / update manager. Stdlib-only; see __main__.py."""
+
+# Version of the CLI contract consumed by external frontends (desktop
+# launcher, install scripts). Bump only on breaking changes to command
+# names/flags or --json output shapes.
+CLI_CONTRACT_VERSION = 1

manager/__main__.py

@@ -0,0 +1,249 @@
+"""AI Toolkit manager CLI.
+
+Runs with any Python >= 3.8 and no dependencies, so it works before the
+training environment exists. This is the single entry point every installer
+frontend (shell scripts, the desktop launcher, the web UI) shells out to.
+
+    python3 -m manager install          first-time environment setup
+    python3 -m manager check [--json]   is an update / dep sync needed?
+    python3 -m manager update           git pull + dependency sync + migrations
+    python3 -m manager sync             dependency sync only (no git pull)
+    python3 -m manager launch           start the web UI
+    python3 -m manager detect [--json]  show detected hardware
+    python3 -m manager doctor           full environment diagnostics
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+# allow `python manager/__main__.py` as well as `python -m manager`
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from manager import detect as detect_mod
+from manager import env, gitops, launch, spec as spec_mod, util
+from manager.util import die, info, ok, print_json, warn
+
+
+def _resolve_spec(args):
+    detection = detect_mod.detect()
+    try:
+        return detection, spec_mod.build_spec(
+            detection, allow_cpu=getattr(args, "cpu", False)
+        )
+    except RuntimeError as e:
+        die(str(e))
+
+
+def cmd_detect(args):
+    detection = detect_mod.detect()
+    try:
+        s = spec_mod.build_spec(detection, allow_cpu=True)
+        detection["spec"] = s.as_dict()
+    except RuntimeError as e:
+        detection["spec_error"] = str(e)
+    if args.json:
+        print_json(detection)
+    else:
+        backend = detection.get("spec", {}).get("backend", "unknown")
+        info("os=%s arch=%s backend=%s" % (detection["os"], detection["arch"], backend))
+        if detection["nvidia"]:
+            for gpu in detection["nvidia"]["gpus"]:
+                info("gpu: %s (%s)" % (gpu["name"], gpu["memory"]))
+
+
+def cmd_install(args):
+    detection, s = _resolve_spec(args)
+    env.sync(s, detection, dry_run=args.dry_run, force=args.force)
+    if not args.dry_run:
+        ok("Install complete. Start the UI with: python3 -m manager launch")
+
+
+def cmd_sync(args):
+    detection, s = _resolve_spec(args)
+    env.sync(s, detection, dry_run=args.dry_run, force=args.force)
+
+
+def cmd_check(args):
+    _, s = _resolve_spec(args)
+    fetched = gitops.fetch()
+    behind = gitops.behind_count()
+    data = {
+        "version": _toolkit_version(),
+        "branch": gitops.current_branch(),
+        "commit": gitops.current_commit(),
+        "remote_commit": gitops.remote_commit(),
+        "dirty": gitops.is_dirty(),
+        "fetch_ok": fetched,
+        "behind": behind,
+        "incoming": gitops.incoming_log(),
+        "venv": env.venv_exists(),
+        "deps_in_sync": env.venv_exists()
+        and env.torch_matches(s)
+        and env.requirements_in_sync(s),
+        "backend": s.backend,
+    }
+    data["update_available"] = bool(behind) or not data["deps_in_sync"]
+    if args.json:
+        print_json(data)
+        return
+    info("AI Toolkit %s (%s @ %s)" % (data["version"], data["branch"], data["commit"]))
+    if not fetched:
+        warn("Could not reach the remote (offline?) — update status may be stale.")
+    if behind:
+        info("Update available: %d new commit(s)." % behind)
+        for line in data["incoming"]:
+            print("    " + line)
+    elif behind == 0:
+        ok("Code is up to date.")
+    if not data["deps_in_sync"]:
+        warn("Dependencies are out of sync. Run: python3 -m manager sync")
+    elif behind == 0:
+        ok("Dependencies are in sync.")
+
+
+def cmd_update(args):
+    """git pull (never destructive) + dependency sync.
+
+    Local work is sacred: a dirty tree either aborts (default), or with
+    --auto is skipped with a warning so run scripts can continue to launch.
+    We never reset/clean; even a forced pull is --ff-only, which git itself
+    aborts rather than overwriting local changes.
+    """
+    auto = getattr(args, "auto", False)
+    skip_pull = False
+
+    if gitops.is_dirty() and not args.force:
+        if auto:
+            warn(
+                "Local changes detected — skipping the code update to protect "
+                "your work. Commit or stash your changes to receive updates."
+            )
+            skip_pull = True
+        else:
+            die(
+                "You have local changes to tracked files. Commit or stash them, "
+                "or re-run with --force to attempt the update anyway (git will "
+                "still refuse rather than overwrite your changes)."
+            )
+
+    if not skip_pull and not gitops.fetch():
+        if auto:
+            warn("Could not reach the git remote — skipping the update check.")
+            skip_pull = True
+        else:
+            die("Could not reach the git remote. Check your network and try again.")
+
+    if not skip_pull:
+        behind = gitops.behind_count()
+        if behind is None:
+            warn("Current branch has no upstream; skipping git pull.")
+        elif behind == 0:
+            ok("Code already up to date.")
+        else:
+            info("Pulling %d new commit(s)..." % behind)
+            gitops.pull_ff()
+            ok("Code updated to %s." % gitops.current_commit())
+            # Re-exec so the freshly pulled manager code runs its own dependency
+            # sync and migrations (the in-memory copy of this module is stale now).
+            cmd = [sys.executable, "-m", "manager", "sync"]
+            if args.dry_run:
+                cmd.append("--dry-run")
+            sys.exit(subprocess.call(cmd, cwd=util.REPO_ROOT))
+    # nothing was pulled — safe to sync with the code already loaded
+    detection, s = _resolve_spec(args)
+    env.sync(s, detection, dry_run=args.dry_run)
+
+
+def cmd_launch(args):
+    sys.exit(launch.launch_ui(open_browser=not args.no_browser))
+
+
+def cmd_doctor(args):
+    from manager import doctor
+
+    doctor.run_doctor()
+
+
+def cmd_version(args):
+    print(_toolkit_version())
+
+
+def _toolkit_version():
+    version = {}
+    try:
+        with open(os.path.join(util.REPO_ROOT, "version.py")) as f:
+            exec(f.read(), version)
+        return version.get("VERSION", "unknown")
+    except OSError:
+        return "unknown"
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        prog="manager", description="AI Toolkit install / update manager"
+    )
+    sub = parser.add_subparsers(dest="command")
+
+    def add(name, fn, **kwargs):
+        p = sub.add_parser(name, **kwargs)
+        p.set_defaults(fn=fn)
+        return p
+
+    p = add("detect", cmd_detect, help="show detected hardware and env spec")
+    p.add_argument("--json", action="store_true")
+
+    for name, fn, help_text in (
+        ("install", cmd_install, "first-time environment setup"),
+        ("sync", cmd_sync, "sync dependencies for the current checkout"),
+    ):
+        p = add(name, fn, help=help_text)
+        p.add_argument("--cpu", action="store_true", help="allow CPU-only install")
+        p.add_argument("--dry-run", action="store_true")
+        p.add_argument(
+            "--force",
+            action="store_true",
+            help="reinstall requirements even if in sync",
+        )
+
+    p = add("check", cmd_check, help="check for updates (use --json for machines)")
+    p.add_argument("--json", action="store_true")
+    p.add_argument("--cpu", action="store_true", help=argparse.SUPPRESS)
+
+    p = add("update", cmd_update, help="git pull + dependency sync + migrations")
+    p.add_argument("--cpu", action="store_true", help="allow CPU-only install")
+    p.add_argument("--dry-run", action="store_true")
+    p.add_argument(
+        "--force", action="store_true", help="update even with local changes"
+    )
+    p.add_argument(
+        "--auto",
+        action="store_true",
+        help="unattended mode (run scripts): on local changes or an unreachable "
+        "remote, warn and skip the code update instead of failing; deps still sync",
+    )
+
+    p = add("launch", cmd_launch, help="start the web UI")
+    p.add_argument(
+        "--no-browser",
+        action="store_true",
+        help="do not open a browser when the UI is ready",
+    )
+    add("doctor", cmd_doctor, help="diagnose the environment")
+    add("version", cmd_version, help="print the toolkit version")
+
+    args = parser.parse_args(argv)
+    if not getattr(args, "command", None):
+        parser.print_help()
+        return 1
+    util.set_json_mode(bool(getattr(args, "json", False)))
+    try:
+        args.fn(args)
+    except KeyboardInterrupt:
+        return 130
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

manager/detect.py

@@ -0,0 +1,110 @@
+"""Hardware / platform detection. Stdlib only, safe to run anywhere."""
+
+import os
+import platform
+import re
+import subprocess
+
+from .util import which, IS_WINDOWS, IS_MAC, IS_LINUX
+
+
+def _run_quiet(cmd):
+    try:
+        out = subprocess.run(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=15
+        )
+        if out.returncode != 0:
+            return None
+        return out.stdout.decode("utf-8", errors="replace")
+    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
+        return None
+
+
+def detect_nvidia():
+    """Returns dict with gpus + driver info, or None if no working nvidia-smi."""
+    smi = which("nvidia-smi")
+    if not smi:
+        return None
+    fields = "name,memory.total,driver_version,compute_cap"
+    csv = _run_quiet([smi, "--query-gpu=" + fields, "--format=csv,noheader"])
+    if not csv:
+        # older drivers don't know the compute_cap field
+        csv = _run_quiet(
+            [
+                smi,
+                "--query-gpu=name,memory.total,driver_version",
+                "--format=csv,noheader",
+            ]
+        )
+    if not csv:
+        return None
+    gpus = []
+    driver = None
+    for line in csv.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) >= 3:
+            gpu = {"name": parts[0], "memory": parts[1]}
+            if len(parts) >= 4:
+                gpu["compute_cap"] = parts[3]
+            gpus.append(gpu)
+            driver = parts[2]
+    if not gpus:
+        return None
+    # Max CUDA version the driver supports only appears in the banner output
+    banner = _run_quiet([smi]) or ""
+    m = re.search(r"CUDA Version:\s*([0-9]+\.[0-9]+)", banner)
+    cuda_version = m.group(1) if m else None
+    return {"gpus": gpus, "driver": driver, "cuda_version": cuda_version}
+
+
+def detect_rocm():
+    """Returns dict if an AMD ROCm stack is present, else None."""
+    if not IS_LINUX:
+        return None
+    smi = which("rocm-smi")
+    if not smi and not os.path.isdir("/opt/rocm"):
+        return None
+    gpus = []
+    out = _run_quiet([smi, "--showproductname"]) if smi else None
+    if out:
+        for m in re.finditer(r"Card [Ss]eries:\s*(.+)", out):
+            gpus.append({"name": m.group(1).strip()})
+    return {"gpus": gpus}
+
+
+def detect():
+    """Full platform detection. Returns a plain dict (json-serializable)."""
+    system = platform.system()  # Linux / Darwin / Windows
+    arch = platform.machine().lower()  # x86_64 / amd64 / arm64 / aarch64
+    if arch == "amd64":
+        arch = "x86_64"
+    if arch == "arm64" and not IS_MAC:
+        arch = "aarch64"
+
+    result = {
+        "os": {"Linux": "linux", "Darwin": "mac", "Windows": "windows"}.get(
+            system, system.lower()
+        ),
+        "arch": arch,
+        "python": platform.python_version(),
+        "nvidia": None,
+        "rocm": None,
+        "backend": "cpu",
+    }
+
+    nvidia = detect_nvidia()
+    if nvidia:
+        result["nvidia"] = nvidia
+        result["backend"] = "cuda"
+    else:
+        rocm = detect_rocm()
+        if rocm:
+            result["rocm"] = rocm
+            result["backend"] = "rocm"
+
+    if IS_MAC:
+        result["backend"] = "mps" if arch == "arm64" else "cpu"
+
+    # DGX OS / Grace (GB10, DGX Spark): NVIDIA GPU on aarch64 Linux
+    result["is_dgx"] = bool(result["os"] == "linux" and arch == "aarch64" and nvidia)
+    return result

manager/doctor.py

@@ -0,0 +1,127 @@
+"""Environment diagnostics: `python -m manager doctor`."""
+
+import os
+import shutil
+import subprocess
+import sys
+
+from . import detect as detect_mod
+from . import env, ffmpeg, gitops, nodejs
+from .util import REPO_ROOT, clean_env, find_uv, venv_dir, venv_python
+
+
+def _check(label, passed, detail=""):
+    if sys.stdout.isatty():
+        mark = "\033[32mOK\033[0m " if passed else "\033[31mFAIL\033[0m"
+    else:
+        mark = "OK  " if passed else "FAIL"
+    print("  [%s] %-18s %s" % (mark, label, detail))
+    return passed
+
+
+def run_doctor():
+    print("AI Toolkit doctor\n")
+    d = detect_mod.detect()
+
+    from . import gitwin
+
+    _check("os / arch", True, "%s %s" % (d["os"], d["arch"]))
+    git = gitwin.find_git()
+    _check(
+        "git",
+        git is not None,
+        git or "not found (manager sync installs a local copy on Windows)",
+    )
+    uv = find_uv()
+    _check("uv", True, uv or "not found (optional, recommended)")
+
+    if d["nvidia"]:
+        names = ", ".join(g["name"] for g in d["nvidia"]["gpus"])
+        _check(
+            "gpu",
+            True,
+            "%s (driver %s, CUDA %s)"
+            % (names, d["nvidia"]["driver"], d["nvidia"]["cuda_version"]),
+        )
+    elif d["rocm"]:
+        _check("gpu", True, "AMD ROCm (experimental)")
+    elif d["backend"] == "mps":
+        _check("gpu", True, "Apple Silicon (MPS)")
+    else:
+        _check("gpu", False, "no supported GPU detected")
+
+    has_venv = env.venv_exists()
+    _check("venv", has_venv, venv_dir() if has_venv else "not created yet")
+    if has_venv:
+        torch = env.installed_torch()
+        _check("torch", torch is not None, torch or "not installed")
+        if torch and d["backend"] == "cuda":
+            try:
+                out = subprocess.run(
+                    [
+                        venv_python(),
+                        "-c",
+                        "import torch; print(torch.cuda.is_available())",
+                    ],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.DEVNULL,
+                    timeout=120,
+                    env=clean_env(),
+                )
+                avail = out.stdout.decode().strip() == "True"
+                _check(
+                    "torch sees gpu",
+                    avail,
+                    "" if avail else "torch.cuda.is_available() is False",
+                )
+            except (OSError, subprocess.TimeoutExpired):
+                _check("torch sees gpu", False, "could not query")
+
+    node_exe, node_major = nodejs.have_usable_node()
+    if node_exe:
+        _check("node", True, "%s (v%s)" % (node_exe, node_major))
+    else:
+        _check(
+            "node",
+            False,
+            "none >= %d found (manager sync installs a local copy)"
+            % nodejs.MIN_NODE_MAJOR,
+        )
+
+    if os.path.isfile(ffmpeg.ffmpeg_exe()):
+        # run it with the launch env so missing shared libs are caught
+        from . import launch
+
+        try:
+            out = subprocess.run(
+                [ffmpeg.ffmpeg_exe(), "-version"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.DEVNULL,
+                timeout=30,
+                env=launch.build_env(),
+            )
+            works = out.returncode == 0
+            detail = (
+                out.stdout.decode().splitlines()[0]
+                if works
+                else "installed but fails to run"
+            )
+        except (OSError, subprocess.TimeoutExpired):
+            works, detail = False, "installed but fails to run"
+        _check("ffmpeg (local)", works, detail)
+    else:
+        _check("ffmpeg (local)", False, "not installed (manager sync installs it)")
+
+    try:
+        free_gb = shutil.disk_usage(REPO_ROOT).free / (1024**3)
+        _check("disk space", free_gb > 30, "%.0f GB free" % free_gb)
+    except OSError:
+        pass
+
+    branch = gitops.current_branch()
+    _check(
+        "git checkout",
+        True,
+        "%s @ %s%s"
+        % (branch, gitops.current_commit(), " (dirty)" if gitops.is_dirty() else ""),
+    )

manager/env.py

@@ -0,0 +1,454 @@
+"""Python environment provisioning and dependency sync.
+
+Strategy:
+- If a venv already exists (.venv or venv), use it.
+- Otherwise create one: prefer uv (downloads the exact Python version needed),
+  fall back to the running Python's venv module if it is new enough.
+- Installs go through `uv pip` when uv is available (much faster), else pip.
+
+State (torch backend, requirements hash, applied migrations) is stored inside
+the venv so a deleted venv means a clean slate — which is correct.
+"""
+
+import json
+import os
+import subprocess
+import sys
+
+from .util import (
+    REPO_ROOT,
+    IS_WINDOWS,
+    clean_env,
+    die,
+    file_hash,
+    find_uv,
+    info,
+    ok,
+    run,
+    venv_dir,
+    venv_python,
+    warn,
+)
+
+STATE_FILE = "aitk_manager_state.json"
+
+MIN_SYSTEM_PYTHON = (3, 10)
+
+
+# ---------------------------------------------------------------- state
+
+
+def state_path(venv=None):
+    return os.path.join(venv or venv_dir(), STATE_FILE)
+
+
+def load_state():
+    try:
+        with open(state_path(), "r") as f:
+            return json.load(f)
+    except (OSError, ValueError):
+        return {}
+
+
+def save_state(state):
+    with open(state_path(), "w") as f:
+        json.dump(state, f, indent=2)
+
+
+# ---------------------------------------------------------------- venv
+
+
+def venv_exists():
+    return os.path.isfile(venv_python())
+
+
+def ensure_venv(spec, dry_run=False):
+    """Create the venv if missing. Returns path to the venv python."""
+    if venv_exists():
+        return venv_python()
+
+    target = venv_dir()
+    uv = find_uv()
+    if dry_run:
+        info(
+            "[dry-run] would create venv at %s (python %s, via %s)"
+            % (target, spec.python_version, "uv" if uv else "venv module")
+        )
+        return venv_python(target)
+
+    if uv:
+        info("Creating venv with uv (python %s) at %s" % (spec.python_version, target))
+        run(
+            [uv, "venv", target, "--python", spec.python_version, "--seed"],
+            env=clean_env(),
+        )
+    else:
+        if sys.version_info < MIN_SYSTEM_PYTHON:
+            die(
+                "Python %d.%d is too old (need >= %d.%d) and uv is not installed.\n"
+                "Install uv (https://docs.astral.sh/uv/) or a newer Python, then re-run."
+                % (sys.version_info[:2] + MIN_SYSTEM_PYTHON)
+            )
+        pyver = "%d.%d" % sys.version_info[:2]
+        if pyver != spec.python_version:
+            warn(
+                "Recommended Python is %s but using system Python %s "
+                "(install uv to get the exact version automatically)."
+                % (spec.python_version, pyver)
+            )
+        info("Creating venv at %s" % target)
+        run([sys.executable, "-m", "venv", target])
+    ok("Virtual environment ready.")
+    return venv_python(target)
+
+
+def _pip_install(args, dry_run=False, upgrade=False, check=True):
+    """Install into the venv, via uv pip if available. Returns exit code."""
+    uv = find_uv()
+    if uv:
+        cmd = [uv, "pip", "install", "--python", venv_python()]
+    else:
+        cmd = [venv_python(), "-m", "pip", "install"]
+    if upgrade:
+        cmd.append("--upgrade")
+    cmd += args
+    if dry_run:
+        info("[dry-run] would run: %s" % " ".join(cmd))
+        return 0
+    code, _ = run(cmd, stream=True, env=clean_env(), check=check)
+    return code
+
+
+def _pip_uninstall(packages, dry_run=False):
+    uv = find_uv()
+    if uv:
+        cmd = [uv, "pip", "uninstall", "--python", venv_python()] + packages
+    else:
+        cmd = [venv_python(), "-m", "pip", "uninstall", "-y"] + packages
+    if dry_run:
+        info("[dry-run] would run: %s" % " ".join(cmd))
+        return
+    # non-fatal: the package may simply not be installed yet
+    run(cmd, check=False, env=clean_env())
+
+
+def venv_python_version():
+    """'3.12' etc. from the venv interpreter, or None."""
+    if not venv_exists():
+        return None
+    try:
+        out = subprocess.run(
+            [venv_python(), "-c", "import sys; print('%d.%d' % sys.version_info[:2])"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            timeout=30,
+            env=clean_env(),
+        )
+        if out.returncode != 0:
+            return None
+        return out.stdout.decode().strip() or None
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+
+
+# ---------------------------------------------------------------- torch
+
+
+def installed_torch():
+    """Returns torch.__version__ from the venv (e.g. '2.9.1+cu128'), or None."""
+    if not venv_exists():
+        return None
+    try:
+        out = subprocess.run(
+            [venv_python(), "-c", "import torch; print(torch.__version__)"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            timeout=120,
+            env=clean_env(),
+        )
+        if out.returncode != 0:
+            return None
+        return out.stdout.decode().strip() or None
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+
+
+def torch_matches(spec):
+    current = installed_torch()
+    if current is None:
+        return False
+    want = spec.torch_packages["torch"]
+    # local version tag carries the backend: "2.9.1+cu128"
+    if "+" in current:
+        version, local = current.split("+", 1)
+        return version == want and local == spec.backend
+    # PyPI wheels (mac) have no local tag
+    return current == want and spec.backend in ("mps", "cpu")
+
+
+def ensure_torch(spec, dry_run=False):
+    if torch_matches(spec):
+        ok(
+            "PyTorch %s (%s) already installed."
+            % (spec.torch_packages["torch"], spec.backend)
+        )
+        return False
+    current = installed_torch()
+    if current:
+        info(
+            "PyTorch %s installed, need %s (%s) — reinstalling."
+            % (current, spec.torch_packages["torch"], spec.backend)
+        )
+    else:
+        info(
+            "Installing PyTorch %s (%s)..."
+            % (spec.torch_packages["torch"], spec.backend)
+        )
+    _pip_install(spec.torch_args(), dry_run=dry_run)
+    return True
+
+
+# ---------------------------------------------------------------- requirements
+
+
+def requirements_hash(spec):
+    """Hash of every requirements file plus the spec itself."""
+    req_files = [
+        os.path.join(REPO_ROOT, f)
+        for f in os.listdir(REPO_ROOT)
+        if f.startswith("requirements") and f.endswith(".txt")
+    ]
+    req_files.append(os.path.join(REPO_ROOT, "dgx_requirements.txt"))
+    base = file_hash(req_files)
+    import hashlib
+
+    h = hashlib.sha256()
+    h.update(base.encode())
+    h.update(json.dumps(spec.as_dict(), sort_keys=True).encode())
+    return h.hexdigest()
+
+
+def requirements_in_sync(spec):
+    if not venv_exists():
+        return False
+    return load_state().get("req_hash") == requirements_hash(spec)
+
+
+def _git_pinned_packages(spec):
+    """{package_name: full git+ requirement line} from the requirements files.
+
+    pip skips reinstalling a git pin whose version number didn't change even
+    when the commit hash did, so pins whose URL changed since the last sync
+    get uninstalled first to force the new commit (the trick the community
+    Windows installer uses for diffusers).
+    """
+    pins = {}
+    seen_files = set()
+
+    def scan(path):
+        if path in seen_files or not os.path.isfile(path):
+            return
+        seen_files.add(path)
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith("-r "):
+                    scan(os.path.join(os.path.dirname(path), line[3:].strip()))
+                elif "git+" in line and not line.startswith("#"):
+                    # e.g. git+https://github.com/huggingface/diffusers.git@<sha>
+                    tail = line.split("/")[-1]
+                    name = tail.split(".git")[0].split("@")[0]
+                    if name:
+                        pins[name] = line
+
+    scan(spec.requirements_path())
+    return pins
+
+
+def _stale_git_pins(spec):
+    """Git-pinned packages whose pin (commit) changed since the last sync."""
+    current = _git_pinned_packages(spec)
+    state = load_state()
+    stored = state.get("git_pins")
+    if stored is None:
+        # no record of what's installed (pre-tracking env): if deps were ever
+        # installed here, play it safe and force-reinstall all git pins once
+        return list(current) if state.get("req_hash") else []
+    return [name for name, line in current.items() if stored.get(name) != line]
+
+
+def _optional_import_name(pkg):
+    """Importable module name for an optional package spec or wheel URL."""
+    if "://" in pkg:
+        name = os.path.basename(pkg).split("-")[0]
+    else:
+        name = pkg
+        for sep in ("==", ">=", "<=", "<", ">", "["):
+            name = name.split(sep)[0]
+    return name.strip().replace("-", "_")
+
+
+def _venv_import_ok(module_name):
+    try:
+        out = subprocess.run(
+            [venv_python(), "-c", "import %s" % module_name],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=180,
+            env=clean_env(),
+        )
+        return out.returncode == 0
+    except (OSError, subprocess.TimeoutExpired):
+        return False
+
+
+def _filter_extras(extras):
+    """Drop wheel URLs whose cpXY tag doesn't match the venv python."""
+    pyver = venv_python_version()
+    cp_tag = "cp" + pyver.replace(".", "") if pyver else None
+    kept = []
+    for pkg in extras:
+        if "cp3" in pkg and cp_tag and cp_tag not in pkg:
+            warn(
+                "Skipping %s (built for a different python than venv %s)."
+                % (os.path.basename(pkg), pyver)
+            )
+            continue
+        kept.append(pkg)
+    return kept
+
+
+def ensure_requirements(spec, dry_run=False, force=False):
+    if not force and requirements_in_sync(spec):
+        ok("Requirements already in sync.")
+        return False
+    # force git-pinned deps (diffusers) onto a newly pinned commit — pip won't
+    # reinstall them on its own because the version number stays the same
+    stale = _stale_git_pins(spec)
+    if stale:
+        info("Git pin changed — reinstalling: %s" % ", ".join(stale))
+        _pip_uninstall(stale, dry_run=dry_run)
+    info("Installing requirements from %s..." % spec.requirements_file)
+    _pip_install(["-r", spec.requirements_path()], dry_run=dry_run)
+    find_links = []
+    for url in spec.find_links:
+        find_links += ["--find-links", url]
+    extras = _filter_extras(spec.extra_packages)
+    if extras:
+        info("Installing platform extras...")
+        _pip_install(extras + find_links, dry_run=dry_run, upgrade=True)
+    # accelerators (flash-attn, NATTEN, ...): install one-by-one, warn on
+    # failure — training works without them, so never fail the whole install
+    for pkg in _filter_extras(spec.optional_packages):
+        label = os.path.basename(pkg) if "://" in pkg else pkg
+        info("Installing optional accelerator: %s" % label)
+        code = _pip_install(
+            [pkg] + find_links, dry_run=dry_run, upgrade=True, check=False
+        )
+        if code != 0:
+            warn("Optional package failed to install (continuing): %s" % label)
+            continue
+        if dry_run:
+            continue
+        # prebuilt accelerator wheels are sometimes built against a torch
+        # nightly and fail to load against the release ABI — verify the
+        # import and roll back rather than leaving a broken wheel installed
+        name = _optional_import_name(pkg)
+        if not _venv_import_ok(name):
+            warn(
+                "%s installed but fails to import against this torch build — "
+                "removing it (training falls back to native attention)." % name
+            )
+            _pip_uninstall([name])
+    if not dry_run:
+        state = load_state()
+        state["req_hash"] = requirements_hash(spec)
+        state["backend"] = spec.backend
+        state["git_pins"] = _git_pinned_packages(spec)
+        save_state(state)
+    return True
+
+
+# ---------------------------------------------------------------- sitecustomize
+
+
+def write_sitecustomize(dry_run=False):
+    """Drop a sitecustomize.py into the venv that exposes the local ffmpeg.
+
+    sitecustomize is imported automatically at interpreter startup, so ANY use
+    of the venv python (UI-spawned training jobs, run.py from a terminal) gets
+    .ffmpeg/bin on PATH — and on Windows, os.add_dll_directory so torchcodec
+    finds the FFmpeg DLLs.
+    """
+    from . import ffmpeg
+
+    if not venv_exists():
+        return
+    try:
+        out = subprocess.run(
+            [
+                venv_python(),
+                "-c",
+                "import sysconfig; print(sysconfig.get_paths()['purelib'])",
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            timeout=30,
+            env=clean_env(),
+        )
+        site_packages = out.stdout.decode().strip()
+    except (OSError, subprocess.TimeoutExpired):
+        site_packages = ""
+    if not site_packages or not os.path.isdir(site_packages):
+        warn("Could not locate venv site-packages — skipping sitecustomize.")
+        return
+    target = os.path.join(site_packages, "sitecustomize.py")
+    content = (
+        "# Generated by the AI Toolkit manager (manager/env.py). Do not edit;\n"
+        "# regenerated on every `manager sync`.\n"
+        "import os\n"
+        "_FFMPEG_BIN = %r\n"
+        "_FFMPEG_LIB = %r\n"
+        "if os.path.isdir(_FFMPEG_BIN):\n"
+        "    os.environ['PATH'] = _FFMPEG_BIN + os.pathsep + os.environ.get('PATH', '')\n"
+        "    if hasattr(os, 'add_dll_directory'):\n"
+        "        try:\n"
+        "            os.add_dll_directory(_FFMPEG_BIN)\n"
+        "        except OSError:\n"
+        "            pass\n"
+        "if os.path.isdir(_FFMPEG_LIB):\n"
+        "    # inherited by child processes (the ffmpeg/ffprobe executables\n"
+        "    # need it to find their own shared libs)\n"
+        "    _prev = os.environ.get('LD_LIBRARY_PATH', '')\n"
+        "    if _FFMPEG_LIB not in _prev.split(os.pathsep):\n"
+        "        os.environ['LD_LIBRARY_PATH'] = (\n"
+        "            _FFMPEG_LIB + ((os.pathsep + _prev) if _prev else '')\n"
+        "        )\n"
+    ) % (ffmpeg.bin_dir(), ffmpeg.lib_dir())
+    if dry_run:
+        info("[dry-run] would write %s" % target)
+        return
+    with open(target, "w") as f:
+        f.write(content)
+
+
+# ---------------------------------------------------------------- sync
+
+
+def sync(spec, detection, dry_run=False, force=False):
+    """Bring the environment fully up to date for this checkout."""
+    from . import ffmpeg, gitwin, migrations, nodejs, uvbin
+
+    for note in spec.notes:
+        warn(note)
+    uvbin.ensure_uv(dry_run=dry_run)
+    gitwin.ensure_git(dry_run=dry_run)
+    ensure_venv(spec, dry_run=dry_run)
+    changed_torch = ensure_torch(spec, dry_run=dry_run)
+    # a torch reinstall can clobber pinned deps; force req pass afterwards
+    ensure_requirements(spec, dry_run=dry_run, force=force or changed_torch)
+    ffmpeg.ensure_ffmpeg(detection, dry_run=dry_run)
+    nodejs.ensure_node(detection, dry_run=dry_run)
+    write_sitecustomize(dry_run=dry_run)
+    migrations.run_pending(dry_run=dry_run)
+    ok("Environment is up to date.")

manager/ffmpeg.py

@@ -0,0 +1,165 @@
+"""Local (never global) FFmpeg provisioning into <repo>/.ffmpeg/.
+
+Why: the toolkit shells out to ffmpeg/ffprobe for video work, and torchcodec
+dlopens the FFmpeg *shared libraries* at runtime. Installing FFmpeg system-wide
+(apt/winget/brew) is exactly what we want to avoid, so we download a portable
+build next to the repo:
+
+- Linux / Windows: BtbN shared builds (bin/ + lib/ with .so/.dll) — the shared
+  libs are what torchcodec needs. The FFmpeg major version must be one the
+  pinned torchcodec supports.
+- macOS: Martin Riedl static ffmpeg/ffprobe executables (no shared libs
+  published; torchcodec on mac keeps whatever it uses today).
+
+Exposure to the rest of the system:
+- `manager launch` prepends .ffmpeg/bin to PATH (and .ffmpeg/lib to
+  LD_LIBRARY_PATH on Linux) so the UI and every training job it spawns see it.
+- env.py writes a sitecustomize.py into the venv that prepends .ffmpeg/bin to
+  PATH and (on Windows) calls os.add_dll_directory — so torchcodec finds the
+  DLLs in ANY use of the venv python, not just via `manager launch`.
+"""
+
+import os
+import shutil
+import stat
+import tempfile
+
+from .util import (
+    IS_MAC,
+    IS_WINDOWS,
+    REPO_ROOT,
+    download,
+    extract_archive,
+    info,
+    ok,
+    warn,
+)
+
+FFMPEG_DIR = os.path.join(REPO_ROOT, ".ffmpeg")
+
+# FFmpeg 8 on all BtbN platforms — torchcodec 0.15 (pinned in spec.py)
+# supports ffmpeg up to 8. Bump these together with the torchcodec pin.
+_BTBN = "https://github.com/BtbN/FFmpeg-Builds/releases/download/latest/"
+_RIEDL = (
+    "https://ffmpeg.martin-riedl.de/redirect/latest/macos/{arch}/release/{tool}.zip"
+)
+
+_SOURCES = {
+    ("linux", "x86_64"): _BTBN + "ffmpeg-n8.1-latest-linux64-gpl-shared-8.1.tar.xz",
+    ("linux", "aarch64"): _BTBN + "ffmpeg-n8.1-latest-linuxarm64-gpl-shared-8.1.tar.xz",
+    ("windows", "x86_64"): _BTBN + "ffmpeg-n8.1-latest-win64-gpl-shared-8.1.zip",
+}
+
+
+def bin_dir():
+    return os.path.join(FFMPEG_DIR, "bin")
+
+
+def lib_dir():
+    return os.path.join(FFMPEG_DIR, "lib")
+
+
+def ffmpeg_exe():
+    return os.path.join(bin_dir(), "ffmpeg.exe" if IS_WINDOWS else "ffmpeg")
+
+
+def is_installed(source_url):
+    marker = os.path.join(FFMPEG_DIR, ".source")
+    if not os.path.isfile(ffmpeg_exe()) or not os.path.isfile(marker):
+        return False
+    with open(marker) as f:
+        return f.read().strip() == source_url
+
+
+def _mark_installed(source_url):
+    with open(os.path.join(FFMPEG_DIR, ".source"), "w") as f:
+        f.write(source_url)
+
+
+def _install_btbn(url):
+    tmp = tempfile.mkdtemp(prefix="aitk_ffmpeg_")
+    try:
+        archive = os.path.join(tmp, os.path.basename(url))
+        download(url, archive, label="ffmpeg")
+        extract_archive(archive, tmp)
+        # archives contain a single top-level dir with bin/ lib/ include/
+        inner = None
+        for name in os.listdir(tmp):
+            path = os.path.join(tmp, name)
+            if os.path.isdir(path) and os.path.isdir(os.path.join(path, "bin")):
+                inner = path
+                break
+        if inner is None:
+            warn("Unexpected ffmpeg archive layout — skipping ffmpeg install.")
+            return False
+        if os.path.isdir(FFMPEG_DIR):
+            shutil.rmtree(FFMPEG_DIR)
+        shutil.move(inner, FFMPEG_DIR)
+        return True
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
+def _install_mac(detection):
+    arch = "arm64" if detection["arch"] == "arm64" else "amd64"
+    tmp = tempfile.mkdtemp(prefix="aitk_ffmpeg_")
+    try:
+        os.makedirs(bin_dir(), exist_ok=True)
+        for tool in ("ffmpeg", "ffprobe"):
+            url = _RIEDL.format(arch=arch, tool=tool)
+            archive = os.path.join(tmp, tool + ".zip")
+            download(url, archive, label=tool)
+            extract_archive(archive, tmp)
+            src = os.path.join(tmp, tool)
+            if not os.path.isfile(src):
+                warn("Unexpected %s archive layout — skipping." % tool)
+                return False
+            dest = os.path.join(bin_dir(), tool)
+            shutil.move(src, dest)
+            os.chmod(
+                dest, os.stat(dest).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+            )
+        return True
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
+def source_url(detection):
+    if detection["os"] == "mac":
+        arch = "arm64" if detection["arch"] == "arm64" else "amd64"
+        return _RIEDL.format(arch=arch, tool="ffmpeg")
+    return _SOURCES.get((detection["os"], detection["arch"]))
+
+
+def ensure_ffmpeg(detection, dry_run=False):
+    url = source_url(detection)
+    if url is None:
+        warn(
+            "No portable FFmpeg source for %s/%s — skipping local ffmpeg."
+            % (detection["os"], detection["arch"])
+        )
+        return False
+    if is_installed(url):
+        ok("Local FFmpeg already installed (.ffmpeg/).")
+        return False
+    if dry_run:
+        info("[dry-run] would install local FFmpeg from %s into %s" % (url, FFMPEG_DIR))
+        return False
+    installed = (
+        _install_mac(detection) if detection["os"] == "mac" else _install_btbn(url)
+    )
+    if installed:
+        _mark_installed(url)
+        ok("Local FFmpeg installed at %s" % FFMPEG_DIR)
+    return installed
+
+
+def env_additions():
+    """(path_dirs, ld_library_dirs) to prepend when launching anything."""
+    paths = []
+    lib_paths = []
+    if os.path.isdir(bin_dir()):
+        paths.append(bin_dir())
+    if not IS_WINDOWS and not IS_MAC and os.path.isdir(lib_dir()):
+        lib_paths.append(lib_dir())
+    return paths, lib_paths

manager/gitops.py

@@ -0,0 +1,92 @@
+"""Git operations for update checking and pulling. Stdlib only."""
+
+import os
+
+from . import gitwin
+from .util import run, REPO_ROOT, die
+
+
+def _git(args, capture=True, check=True):
+    git = gitwin.find_git()
+    if not git:
+        die(
+            "git was not found. On Windows run `python -m manager sync` to "
+            "install a local copy; elsewhere install git with your package "
+            "manager (or xcode-select --install on macOS)."
+        )
+    # skip LFS payloads — the repo may carry LFS files not needed at runtime
+    env = os.environ.copy()
+    env["GIT_LFS_SKIP_SMUDGE"] = "1"
+    return run([git] + args, cwd=REPO_ROOT, capture=capture, check=check, env=env)
+
+
+def current_branch():
+    _, out = _git(["rev-parse", "--abbrev-ref", "HEAD"])
+    return out
+
+
+def current_commit():
+    _, out = _git(["rev-parse", "--short", "HEAD"])
+    return out
+
+
+def is_dirty():
+    _, out = _git(["status", "--porcelain"])
+    # untracked files don't block updates; modified/staged tracked files do
+    for line in (out or "").splitlines():
+        if not line.startswith("??"):
+            return True
+    return False
+
+
+def fetch():
+    code, _ = _git(["fetch", "--quiet"], capture=False, check=False)
+    return code == 0
+
+
+def upstream():
+    code, out = _git(
+        ["rev-parse", "--abbrev-ref", "--symbolic-full-name", "@{u}"], check=False
+    )
+    if code != 0:
+        return None
+    return out
+
+
+def behind_count():
+    """Commits the local branch is behind its upstream. None if no upstream."""
+    up = upstream()
+    if not up:
+        return None
+    code, out = _git(["rev-list", "--count", "HEAD..@{u}"], check=False)
+    if code != 0 or out is None:
+        return None
+    try:
+        return int(out)
+    except ValueError:
+        return None
+
+
+def remote_commit():
+    code, out = _git(["rev-parse", "--short", "@{u}"], check=False)
+    return out if code == 0 else None
+
+
+def incoming_log(limit=15):
+    code, out = _git(
+        ["log", "--oneline", "HEAD..@{u}", "--max-count=%d" % limit], check=False
+    )
+    if code != 0 or not out:
+        return []
+    return out.splitlines()
+
+
+def pull_ff():
+    """Fast-forward pull. Dies with guidance on failure."""
+    code, _ = _git(["pull", "--ff-only"], capture=False, check=False)
+    if code != 0:
+        die(
+            "git pull --ff-only failed. Your branch has local commits or "
+            "conflicts with the remote. Resolve manually (git stash / git rebase) "
+            "and re-run the update."
+        )

manager/gitwin.py

@@ -0,0 +1,65 @@
+"""Local (never global) Git provisioning for Windows into <repo>/.mingit/.
+
+MinGit is the official minimal, portable Git for Windows — a plain zip meant
+for embedding in tools (no installer, no registry, no PATH changes). On
+Windows, if the system has no git, `manager sync` drops one here so updates
+keep working from any terminal. The first clone (before this repo exists) is
+handled the same way by the bootstrap layer (install.ps1 / desktop launcher),
+which then moves its MinGit into the fresh checkout as .mingit/.
+
+Linux and macOS have no sane portable git (glibc / Xcode CLT entanglement),
+and git is effectively always available there — so this is Windows-only and
+the manager just errors with install instructions elsewhere.
+"""
+
+import os
+import shutil
+import tempfile
+
+from .util import IS_WINDOWS, REPO_ROOT, download, extract_archive, info, ok, warn
+
+MINGIT_DIR = os.path.join(REPO_ROOT, ".mingit")
+# Update this pin together with nothing else — it's independent of torch etc.
+MINGIT_URL = (
+    "https://github.com/git-for-windows/git/releases/download/"
+    "v2.55.0.windows.3/MinGit-2.55.0.3-64-bit.zip"
+)
+
+
+def local_git_exe():
+    return os.path.join(MINGIT_DIR, "cmd", "git.exe")
+
+
+def find_git():
+    """Path/command for git: repo-local MinGit first (Windows), then system."""
+    if IS_WINDOWS and os.path.isfile(local_git_exe()):
+        return local_git_exe()
+    return shutil.which("git")
+
+
+def ensure_git(dry_run=False):
+    """Windows-only: provision .mingit/ when the system has no git."""
+    if not IS_WINDOWS:
+        return False
+    if find_git():
+        return False
+    if dry_run:
+        info("[dry-run] would install MinGit into %s" % MINGIT_DIR)
+        return False
+    tmp = tempfile.mkdtemp(prefix="aitk_mingit_")
+    try:
+        archive = os.path.join(tmp, "mingit.zip")
+        download(MINGIT_URL, archive, label="MinGit")
+        # MinGit zips have no top-level folder: cmd/, mingw64/, etc at the root
+        extracted = os.path.join(tmp, "mingit")
+        extract_archive(archive, extracted)
+        if not os.path.isfile(os.path.join(extracted, "cmd", "git.exe")):
+            warn("Unexpected MinGit archive layout — skipping local git install.")
+            return False
+        if os.path.isdir(MINGIT_DIR):
+            shutil.rmtree(MINGIT_DIR)
+        shutil.move(extracted, MINGIT_DIR)
+        ok("Local Git (MinGit) installed at %s" % MINGIT_DIR)
+        return True
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)

manager/launch.py

@@ -0,0 +1,116 @@
+"""Launch the AI Toolkit web UI (ui/ -> npm run build_and_start)."""
+
+import os
+import subprocess
+import threading
+
+from . import ffmpeg, nodejs
+from .util import (
+    IS_LINUX,
+    IS_WINDOWS,
+    REPO_ROOT,
+    clean_env,
+    die,
+    info,
+    venv_dir,
+    venv_python,
+)
+
+UI_PORT = 8675
+UI_URL = "http://localhost:%d" % UI_PORT
+BROWSER_POLL_SECONDS = 300
+
+
+def build_env():
+    """Scrubbed env with local node, ffmpeg, and the venv on PATH.
+
+    Everything the UI worker spawns (training jobs) inherits this, so the
+    local ffmpeg/node are visible to the whole process tree.
+    """
+    env = clean_env()
+    path_dirs = []
+    if os.path.isdir(nodejs.node_bin_dir()):
+        path_dirs.append(nodejs.node_bin_dir())
+    ff_paths, ff_libs = ffmpeg.env_additions()
+    path_dirs += ff_paths
+    vbin = os.path.join(venv_dir(), "Scripts" if IS_WINDOWS else "bin")
+    if os.path.isdir(vbin):
+        path_dirs.append(vbin)
+    if path_dirs:
+        env["PATH"] = os.pathsep.join(path_dirs) + os.pathsep + env.get("PATH", "")
+    if ff_libs:
+        env["LD_LIBRARY_PATH"] = os.pathsep.join(
+            ff_libs + [env.get("LD_LIBRARY_PATH", "")]
+        ).rstrip(os.pathsep)
+    return env
+
+
+def _find_npm(env):
+    import shutil
+
+    return shutil.which("npm.cmd" if IS_WINDOWS else "npm", path=env.get("PATH"))
+
+
+def _headless():
+    return IS_LINUX and not (
+        os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")
+    )
+
+
+def _open_browser_when_ready(stop_event):
+    """Poll the UI port in the background; open the browser once it responds."""
+    import time
+    import urllib.request
+    import webbrowser
+
+    waited = 0
+    while not stop_event.is_set() and waited < BROWSER_POLL_SECONDS:
+        try:
+            urllib.request.urlopen(UI_URL, timeout=2).close()
+            webbrowser.open(UI_URL)
+            return
+        except OSError:
+            time.sleep(2)
+            waited += 2
+
+
+def launch_ui(open_browser=True):
+    if not os.path.isfile(venv_python()):
+        die("No Python environment found. Run: python3 -m manager install")
+
+    env = build_env()
+    npm = _find_npm(env)
+    if not npm:
+        die(
+            "Node.js was not found. Run `python3 -m manager sync` to install a "
+            "local copy, or install Node.js >= %d from https://nodejs.org."
+            % nodejs.MIN_NODE_MAJOR
+        )
+    _, major = nodejs.have_usable_node()
+    if major is not None and major < nodejs.MIN_NODE_MAJOR:
+        die(
+            "Node.js v%d found, but >= %d is required. Run `python3 -m manager sync`."
+            % (major, nodejs.MIN_NODE_MAJOR)
+        )
+
+    info("Starting AI Toolkit UI (%s) ..." % UI_URL)
+    stop_event = threading.Event()
+    if open_browser and not _headless():
+        threading.Thread(
+            target=_open_browser_when_ready, args=(stop_event,), daemon=True
+        ).start()
+
+    proc = subprocess.Popen(
+        [npm, "run", "build_and_start"], cwd=os.path.join(REPO_ROOT, "ui"), env=env
+    )
+    try:
+        return proc.wait()
+    except KeyboardInterrupt:
+        proc.terminate()
+        try:
+            proc.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+        return 130
+    finally:
+        stop_event.set()

manager/migrations.py

@@ -0,0 +1,34 @@
+"""One-time migration steps that run after an update.
+
+Add a migration when an update needs more than a dependency sync (moving
+files, converting configs, clearing caches, etc). Each runs at most once per
+environment; applied ids are recorded in the venv state file.
+
+    def _example(dry_run):
+        ...
+
+    MIGRATIONS = [
+        {"id": "2026-07-example-cache-move", "run": _example},
+    ]
+"""
+
+from .util import info
+from . import env
+
+MIGRATIONS = []
+
+
+def run_pending(dry_run=False):
+    if not MIGRATIONS:
+        return
+    state = env.load_state()
+    applied = set(state.get("migrations", []))
+    for migration in MIGRATIONS:
+        if migration["id"] in applied:
+            continue
+        info("Running migration: %s" % migration["id"])
+        if not dry_run:
+            migration["run"](dry_run)
+            applied.add(migration["id"])
+            state["migrations"] = sorted(applied)
+            env.save_state(state)

manager/nodejs.py

@@ -0,0 +1,121 @@
+"""Local (never global) Node.js provisioning into <repo>/.node/.
+
+A system Node >= 20 is used when present; otherwise an official portable
+build is downloaded next to the repo (same approach run_mac.zsh already uses).
+Nothing is ever installed system-wide.
+"""
+
+import os
+import shutil
+import subprocess
+import tempfile
+
+from .util import (
+    IS_MAC,
+    IS_WINDOWS,
+    REPO_ROOT,
+    download,
+    extract_archive,
+    info,
+    ok,
+    warn,
+    which,
+)
+
+NODE_DIR = os.path.join(REPO_ROOT, ".node")
+# Node 24 is the current LTS line and matches the dgx_instructions.md guidance.
+NODE_VERSION = "24.11.1"
+MIN_NODE_MAJOR = 20
+
+
+def node_bin_dir():
+    # windows zips have node.exe/npm.cmd at the archive root; unix under bin/
+    return NODE_DIR if IS_WINDOWS else os.path.join(NODE_DIR, "bin")
+
+
+def local_node_exe():
+    return os.path.join(node_bin_dir(), "node.exe" if IS_WINDOWS else "node")
+
+
+def _node_major(exe):
+    try:
+        out = subprocess.run(
+            [exe, "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            timeout=15,
+        )
+        if out.returncode != 0:
+            return None
+        text = out.stdout.decode().strip()  # v24.11.1
+        return int(text.lstrip("v").split(".")[0])
+    except (OSError, subprocess.TimeoutExpired, ValueError):
+        return None
+
+
+def _dist_url(detection):
+    arch = detection["arch"]
+    if detection["os"] == "linux":
+        plat = "linux-arm64" if arch == "aarch64" else "linux-x64"
+        ext = "tar.xz"
+    elif detection["os"] == "mac":
+        plat = "darwin-arm64" if arch == "arm64" else "darwin-x64"
+        ext = "tar.gz"
+    elif detection["os"] == "windows":
+        plat = "win-x64"
+        ext = "zip"
+    else:
+        return None, None
+    name = "node-v%s-%s" % (NODE_VERSION, plat)
+    return "https://nodejs.org/dist/v%s/%s.%s" % (NODE_VERSION, name, ext), name
+
+
+def have_usable_node():
+    """(exe, major) for the best available node: local .node/ first, then system."""
+    local = local_node_exe()
+    if os.path.isfile(local):
+        major = _node_major(local)
+        if major is not None and major >= MIN_NODE_MAJOR:
+            return local, major
+    system = which("node")
+    if system:
+        major = _node_major(system)
+        if major is not None and major >= MIN_NODE_MAJOR:
+            return system, major
+    return None, None
+
+
+def ensure_node(detection, dry_run=False):
+    exe, major = have_usable_node()
+    if exe:
+        ok("Node.js v%d found (%s)." % (major, exe))
+        return False
+    url, inner_name = _dist_url(detection)
+    if url is None:
+        warn(
+            "No portable Node.js build for this platform — install Node >= %d manually."
+            % MIN_NODE_MAJOR
+        )
+        return False
+    if dry_run:
+        info(
+            "[dry-run] would install portable Node.js v%s into %s"
+            % (NODE_VERSION, NODE_DIR)
+        )
+        return False
+    tmp = tempfile.mkdtemp(prefix="aitk_node_")
+    try:
+        archive = os.path.join(tmp, os.path.basename(url))
+        download(url, archive, label="node v%s" % NODE_VERSION)
+        extract_archive(archive, tmp)
+        inner = os.path.join(tmp, inner_name)
+        if not os.path.isdir(inner):
+            warn("Unexpected Node.js archive layout — skipping node install.")
+            return False
+        if os.path.isdir(NODE_DIR):
+            shutil.rmtree(NODE_DIR)
+        shutil.move(inner, NODE_DIR)
+        ok("Portable Node.js v%s installed at %s" % (NODE_VERSION, NODE_DIR))
+        return True
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)

manager/spec.py

@@ -0,0 +1,259 @@
+"""Maps detected hardware to an environment spec.
+
+This is the single source of truth for "what does this machine need to run
+this commit of AI Toolkit". One universal torch version across all platforms;
+per-platform accelerator extras (flash-attn, NATTEN, triton) wherever prebuilt
+wheels exist. **Update the pins below together with the README install
+instructions and run_mac.zsh.**
+
+Wheel coverage for the pinned set (verified 2026-07, flash-attn + NATTEN GPU
+kernels smoke-tested on an RTX 5090 / sm120 with torch 2.13.0+cu130):
+- torch 2.13.0: cu126/cu130 wheels for linux x86_64 + aarch64 + windows; PyPI
+  wheels for mac arm64. torchaudio is in maintenance mode — 2.11.0 is the
+  current release and is torch-version-agnostic (no torch dep in metadata).
+- torchcodec 0.15.0: supports torch >= 2.11, wheels on all platforms.
+- flash-attn 2.8.3: prebuilt by mjun0812/flash-attention-prebuild-wheels for
+  {cu126,cu130} x {cp310..cp314} x {linux x86_64, linux aarch64, windows}.
+  NOTE: the torch2.12 linux wheels there were built against a torch nightly
+  and fail to import on 2.12.0 final — the torch2.13 batches (v0.9.47+) are
+  verified good. Re-verify imports whenever bumping torch.
+- NATTEN 0.21.7: prebuilt at whl.natten.org for {cu126,cu130,cu132} x
+  {cp310..cp314} x {linux x86_64, linux aarch64}. No Windows/mac wheels.
+- triton: bundled with torch on Linux (incl. aarch64; torch 2.13 bundles
+  triton 3.7.1); triton-windows 3.7.x matches on Windows; nothing for MPS.
+
+extra_packages are installed AFTER requirements.txt with --upgrade so they can
+override requirement pins (e.g. torchcodec). optional_packages are installed
+one-by-one and only warn on failure (accelerators the training code can live
+without). Wheel URLs containing a cpXY tag are skipped automatically if the
+venv python doesn't match.
+"""
+
+import os
+
+from .util import REPO_ROOT
+
+# ---- version pins (edit these to move the fleet forward) -------------------
+
+TORCH = {"torch": "2.13.0", "torchvision": "0.28.0", "torchaudio": "2.11.0"}
+TORCH_TAG = "2.13"  # as it appears in flash-attn / natten wheel names
+
+TORCHCODEC = "torchcodec==0.15.0"
+TRITON_WINDOWS = "triton-windows>=3.7,<3.8"
+
+NATTEN_VERSION = "0.21.7"
+NATTEN_FIND_LINKS = "https://whl.natten.org"
+
+FLASH_ATTN_VERSION = "2.8.3"
+_FA_BASE = (
+    "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/"
+)
+# (os, arch) -> (release tag, wheel platform tag) — tags are per torch
+# version; these carry the torch2.13 builds
+_FA_BUILDS = {
+    ("linux", "x86_64"): ("v0.9.47", "manylinux_2_24_x86_64.manylinux_2_28_x86_64"),
+    ("linux", "aarch64"): ("v0.9.48", "manylinux_2_34_aarch64"),
+    ("windows", "x86_64"): ("v0.9.52", "win_amd64"),
+}
+
+# helper build tools some sdists need on Windows
+_WIN_HELPERS = ["wheel", "setuptools", "poetry-core", "hf_xet"]
+
+PYTORCH_INDEX = "https://download.pytorch.org/whl/"
+
+
+class EnvSpec(object):
+    def __init__(
+        self,
+        backend,
+        torch_packages,
+        torch_index=None,
+        python_version="3.12",
+        requirements_file="requirements.txt",
+        extra_packages=None,
+        optional_packages=None,
+        find_links=None,
+        notes=None,
+    ):
+        self.backend = backend  # cu130 / cu126 / rocm7.1 / mps / cpu
+        self.torch_packages = torch_packages  # {name: version}
+        self.torch_index = torch_index  # None = PyPI
+        self.python_version = python_version
+        self.requirements_file = requirements_file
+        self.extra_packages = extra_packages or []
+        self.optional_packages = optional_packages or []
+        self.find_links = find_links or []
+        self.notes = notes or []
+
+    def torch_args(self):
+        args = ["%s==%s" % (k, v) for k, v in sorted(self.torch_packages.items())]
+        if self.torch_index:
+            args += ["--index-url", self.torch_index]
+        return args
+
+    def requirements_path(self):
+        return os.path.join(REPO_ROOT, self.requirements_file)
+
+    def as_dict(self):
+        return {
+            "backend": self.backend,
+            "torch_packages": self.torch_packages,
+            "torch_index": self.torch_index,
+            "python_version": self.python_version,
+            "requirements_file": self.requirements_file,
+            "extra_packages": self.extra_packages,
+            "optional_packages": self.optional_packages,
+            "find_links": self.find_links,
+            "notes": self.notes,
+        }
+
+
+def _flash_attn_url(flavor, os_name, arch, python_version):
+    build = _FA_BUILDS.get((os_name, arch))
+    if build is None:
+        return None
+    tag, plat = build
+    cp = "cp" + python_version.replace(".", "")
+    return "%s%s/flash_attn-%s+%storch%s-%s-%s-%s.whl" % (
+        _FA_BASE,
+        tag,
+        FLASH_ATTN_VERSION,
+        flavor,
+        TORCH_TAG,
+        cp,
+        cp,
+        plat,
+    )
+
+
+def _natten_pin(flavor):
+    # natten wheel local tags use the full torch version without dots: torch2120cu130
+    return "natten==%s+torch%s%s" % (
+        NATTEN_VERSION,
+        TORCH["torch"].replace(".", ""),
+        flavor.replace(".", ""),
+    )
+
+
+def _cuda_flavor(detection):
+    """Pick a cuda wheel flavor the installed driver can actually run."""
+    nvidia = detection.get("nvidia") or {}
+    cuda = None
+    if nvidia.get("cuda_version"):
+        try:
+            cuda = tuple(int(x) for x in nvidia["cuda_version"].split("."))
+        except ValueError:
+            cuda = None
+    if cuda is None:
+        # driver present but version unknown — assume current
+        return "cu130", []
+    if cuda >= (13, 0):
+        return "cu130", []
+    caps = [g.get("compute_cap") for g in nvidia.get("gpus", [])]
+    has_blackwell = any(c and float(c) >= 12.0 for c in caps if c)
+    if cuda >= (12, 6):
+        if has_blackwell:
+            raise RuntimeError(
+                "Blackwell GPU detected but the NVIDIA driver only supports "
+                "CUDA %s. Blackwell needs the cu130 build — update your "
+                "driver to 580+ and re-run." % nvidia["cuda_version"]
+            )
+        return "cu126", [
+            "NVIDIA driver only supports CUDA %s — installing cu126 wheels. "
+            "Updating your driver is recommended." % nvidia["cuda_version"]
+        ]
+    raise RuntimeError(
+        "NVIDIA driver only supports CUDA %s, which is too old for the pinned "
+        "torch build. Update your NVIDIA driver, then re-run install."
+        % nvidia["cuda_version"]
+    )
+
+
+def _cuda_spec(detection):
+    os_name = detection["os"]
+    arch = detection["arch"]
+    flavor, notes = _cuda_flavor(detection)
+    python_version = "3.12"
+    requirements = (
+        "dgx_requirements.txt" if detection.get("is_dgx") else "requirements.txt"
+    )
+    if detection.get("is_dgx"):
+        # the old "Python 3.11 on DGX OS" constraint was for conda/system
+        # installs; uv provisions 3.12 and all aarch64 cp312 wheels exist now
+        notes = notes + [
+            "DGX OS / Grace detected: using %s wheels and dgx_requirements.txt."
+            % flavor
+        ]
+
+    extras = [TORCHCODEC]
+    optional = []
+    find_links = []
+
+    fa_url = _flash_attn_url(flavor, os_name, arch, python_version)
+    if fa_url:
+        optional.append(fa_url)
+
+    if os_name == "linux":
+        optional.append(_natten_pin(flavor))
+        find_links.append(NATTEN_FIND_LINKS)
+    elif os_name == "windows":
+        extras = _WIN_HELPERS + extras + [TRITON_WINDOWS]
+        notes = notes + ["NATTEN has no Windows wheels — skipping it."]
+
+    return EnvSpec(
+        flavor,
+        TORCH,
+        torch_index=PYTORCH_INDEX + flavor,
+        python_version=python_version,
+        requirements_file=requirements,
+        extra_packages=extras,
+        optional_packages=optional,
+        find_links=find_links,
+        notes=notes,
+    )
+
+
+def build_spec(detection, allow_cpu=False):
+    """Returns EnvSpec, or raises RuntimeError with a user-facing message."""
+    os_name = detection["os"]
+
+    if os_name == "mac":
+        notes = ["flash-attn / NATTEN / triton are unavailable on macOS."]
+        if detection["backend"] != "mps":
+            notes.append("Intel Mac detected — training will be extremely slow.")
+        return EnvSpec(
+            "mps",
+            TORCH,
+            python_version="3.12",
+            extra_packages=[TORCHCODEC],
+            notes=notes,
+        )
+
+    if detection["backend"] == "cuda":
+        return _cuda_spec(detection)
+
+    if detection["backend"] == "rocm":
+        return EnvSpec(
+            "rocm7.1",
+            TORCH,
+            torch_index=PYTORCH_INDEX + "rocm7.1",
+            extra_packages=[TORCHCODEC],
+            notes=[
+                "AMD ROCm support is experimental and largely untested.",
+                "flash-attn / NATTEN prebuilt wheels are unavailable for ROCm.",
+            ],
+        )
+
+    # CPU fallback
+    if not allow_cpu:
+        raise RuntimeError(
+            "No supported GPU detected (NVIDIA CUDA, AMD ROCm, or Apple Silicon). "
+            "Training on CPU is not practical. Pass --cpu to install anyway."
+        )
+    return EnvSpec(
+        "cpu",
+        TORCH,
+        torch_index=PYTORCH_INDEX + "cpu",
+        extra_packages=[TORCHCODEC],
+        notes=["CPU-only install: training will be impractically slow."],
+    )

manager/util.py

@@ -0,0 +1,230 @@
+"""Shared helpers for the AI Toolkit manager. Stdlib only."""
+
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+IS_WINDOWS = platform.system() == "Windows"
+IS_MAC = platform.system() == "Darwin"
+IS_LINUX = platform.system() == "Linux"
+
+# When --json is used, human output goes to stderr so stdout stays machine-readable
+_json_mode = False
+
+
+def set_json_mode(enabled):
+    global _json_mode
+    _json_mode = enabled
+
+
+def _supports_color(stream):
+    if os.environ.get("NO_COLOR"):
+        return False
+    return hasattr(stream, "isatty") and stream.isatty()
+
+
+def _emit(prefix, msg, color):
+    stream = sys.stderr if _json_mode else sys.stdout
+    if _supports_color(stream):
+        stream.write("\033[%sm%s\033[0m %s\n" % (color, prefix, msg))
+    else:
+        stream.write("%s %s\n" % (prefix, msg))
+    stream.flush()
+
+
+def info(msg):
+    _emit("[*]", msg, "36")
+
+
+def ok(msg):
+    _emit("[+]", msg, "32")
+
+
+def warn(msg):
+    _emit("[!]", msg, "33")
+
+
+def error(msg):
+    _emit("[x]", msg, "31")
+
+
+def die(msg, code=1):
+    error(msg)
+    sys.exit(code)
+
+
+def print_json(data):
+    sys.stdout.write(json.dumps(data, indent=2) + "\n")
+    sys.stdout.flush()
+
+
+def run(cmd, cwd=None, capture=False, check=True, env=None, stream=False):
+    """Run a command. capture=True returns stdout text (stripped).
+
+    stream=True inherits stdio so the user sees live output.
+    Returns (returncode, stdout_or_None).
+    """
+    kwargs = {"cwd": cwd or REPO_ROOT}
+    if env is not None:
+        kwargs["env"] = env
+    if capture:
+        kwargs["stdout"] = subprocess.PIPE
+        kwargs["stderr"] = subprocess.PIPE
+    elif _json_mode and not stream:
+        # keep stdout clean in json mode
+        kwargs["stdout"] = sys.stderr
+
+    try:
+        proc = subprocess.run(cmd, **kwargs)
+    except FileNotFoundError:
+        if check:
+            die("Command not found: %s" % cmd[0])
+        return 127, None
+
+    out = None
+    if capture:
+        out = proc.stdout.decode("utf-8", errors="replace").strip()
+    if check and proc.returncode != 0:
+        detail = ""
+        if capture and proc.stderr:
+            detail = "\n" + proc.stderr.decode("utf-8", errors="replace").strip()
+        die("Command failed (%d): %s%s" % (proc.returncode, " ".join(cmd), detail))
+    return proc.returncode, out
+
+
+def which(name):
+    return shutil.which(name)
+
+
+def find_uv():
+    """Find uv: repo-local .uv/ first, then PATH, then common install dirs."""
+    home = os.path.expanduser("~")
+    candidates = [
+        os.path.join(REPO_ROOT, ".uv", "uv.exe" if IS_WINDOWS else "uv"),
+    ]
+    for c in candidates:
+        if os.path.isfile(c) and os.access(c, os.X_OK):
+            return c
+    uv = shutil.which("uv")
+    if uv:
+        return uv
+    candidates = [
+        os.path.join(home, ".local", "bin", "uv"),
+        os.path.join(home, ".cargo", "bin", "uv"),
+    ]
+    if IS_WINDOWS:
+        local = os.environ.get("LOCALAPPDATA", "")
+        if local:
+            candidates.append(os.path.join(local, "uv", "uv.exe"))
+        candidates.append(os.path.join(home, ".local", "bin", "uv.exe"))
+    for c in candidates:
+        if os.path.isfile(c) and os.access(c, os.X_OK):
+            return c
+    return None
+
+
+def venv_dir():
+    """Existing venv dir (.venv preferred, matching ui/cron/pythonPath.ts), else default target."""
+    for name in (".venv", "venv"):
+        d = os.path.join(REPO_ROOT, name)
+        if os.path.isdir(d):
+            return d
+    return os.path.join(REPO_ROOT, ".venv")
+
+
+def venv_python(venv=None):
+    venv = venv or venv_dir()
+    if IS_WINDOWS:
+        return os.path.join(venv, "Scripts", "python.exe")
+    return os.path.join(venv, "bin", "python3")
+
+
+# Env vars that let a system/conda/pyenv Python leak into our subprocesses.
+# Scrubbed from every python/pip/node invocation (mirrors what the community
+# Windows installer learned the hard way).
+_SCRUB_VARS = (
+    "PYTHONPATH",
+    "PYTHONHOME",
+    "PYTHON",
+    "PYTHONSTARTUP",
+    "PYTHONUSERBASE",
+    "PYTHONEXECUTABLE",
+    "PIP_CONFIG_FILE",
+    "PIP_REQUIRE_VIRTUALENV",
+    "VIRTUAL_ENV",
+    "CONDA_PREFIX",
+    "CONDA_DEFAULT_ENV",
+    "PYENV_ROOT",
+    "PYENV_VERSION",
+)
+
+
+def clean_env(extra=None):
+    """os.environ copy with Python-hijacking vars removed.
+
+    Also points uv's managed-python store into the repo (.uv/python) so
+    interpreter downloads never land outside the checkout.
+    """
+    env = os.environ.copy()
+    for var in _SCRUB_VARS:
+        env.pop(var, None)
+    env.setdefault("UV_PYTHON_INSTALL_DIR", os.path.join(REPO_ROOT, ".uv", "python"))
+    if extra:
+        env.update(extra)
+    return env
+
+
+def download(url, dest, label=None):
+    """Download url to dest (stdlib only), logging progress every ~10%."""
+    import urllib.request
+
+    label = label or os.path.basename(dest)
+    info("Downloading %s ..." % label)
+    last = [-1]
+
+    def hook(blocks, block_size, total):
+        if total <= 0:
+            return
+        pct = min(100, int(blocks * block_size * 100 / total))
+        if pct >= last[0] + 10:
+            last[0] = pct
+            info("  %s: %d%%" % (label, pct))
+
+    tmp = dest + ".part"
+    try:
+        urllib.request.urlretrieve(url, tmp, reporthook=hook)
+    except Exception as e:  # noqa: BLE001 - surface any network failure clearly
+        if os.path.exists(tmp):
+            os.remove(tmp)
+        die("Download failed for %s: %s" % (url, e))
+    os.replace(tmp, dest)
+
+
+def extract_archive(archive, dest_dir):
+    """Extract .zip / .tar.* into dest_dir (created if needed)."""
+    import tarfile
+    import zipfile
+
+    os.makedirs(dest_dir, exist_ok=True)
+    if archive.endswith(".zip"):
+        with zipfile.ZipFile(archive) as z:
+            z.extractall(dest_dir)
+    else:
+        with tarfile.open(archive) as t:
+            t.extractall(dest_dir)
+
+
+def file_hash(paths):
+    import hashlib
+
+    h = hashlib.sha256()
+    for p in sorted(paths):
+        if os.path.isfile(p):
+            with open(p, "rb") as f:
+                h.update(f.read())
+    return h.hexdigest()

manager/uvbin.py

@@ -0,0 +1,78 @@
+"""Local (never system) uv provisioning into <repo>/.uv/.
+
+uv is a single static binary (no deps, no Python needed). Keeping it inside
+the repo means fast installs + exact-version Python provisioning with zero
+footprint outside the checkout. The run scripts also install into this same
+dir (via the astral installer with UV_INSTALL_DIR) for the no-python
+bootstrap case; util.find_uv() checks here first either way.
+"""
+
+import os
+import platform
+import shutil
+import stat
+import tempfile
+
+from .util import IS_MAC, IS_WINDOWS, REPO_ROOT, download, extract_archive, ok, warn
+
+UV_DIR = os.path.join(REPO_ROOT, ".uv")
+
+_RELEASE = "https://github.com/astral-sh/uv/releases/latest/download/"
+
+
+def _asset():
+    arch = platform.machine().lower()
+    if arch == "amd64":
+        arch = "x86_64"
+    if IS_WINDOWS:
+        return "uv-x86_64-pc-windows-msvc.zip"
+    if IS_MAC:
+        return "uv-%s-apple-darwin.tar.gz" % (
+            "aarch64" if arch == "arm64" else "x86_64"
+        )
+    return "uv-%s-unknown-linux-gnu.tar.gz" % (
+        "aarch64" if arch in ("aarch64", "arm64") else "x86_64"
+    )
+
+
+def local_uv_exe():
+    return os.path.join(UV_DIR, "uv.exe" if IS_WINDOWS else "uv")
+
+
+def ensure_uv(dry_run=False):
+    """Download the uv binary into .uv/ if it isn't available anywhere."""
+    from .util import find_uv
+
+    if find_uv():
+        return False
+    if dry_run:
+        from .util import info
+
+        info("[dry-run] would download uv into %s" % UV_DIR)
+        return False
+    tmp = tempfile.mkdtemp(prefix="aitk_uv_")
+    try:
+        asset = _asset()
+        archive = os.path.join(tmp, asset)
+        download(_RELEASE + asset, archive, label="uv")
+        extract_archive(archive, tmp)
+        # zip: uv.exe at root; tarballs: uv-<triple>/uv — search the tree
+        exe_name = "uv.exe" if IS_WINDOWS else "uv"
+        found = None
+        for root, _dirs, files in os.walk(tmp):
+            if exe_name in files:
+                found = os.path.join(root, exe_name)
+                break
+        if not found:
+            warn("Unexpected uv archive layout — continuing without uv.")
+            return False
+        os.makedirs(UV_DIR, exist_ok=True)
+        dest = local_uv_exe()
+        shutil.move(found, dest)
+        os.chmod(
+            dest, os.stat(dest).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+        )
+        ok("uv installed at %s" % dest)
+        return True
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)

requirements_base.txt

@@ -1,4 +1,4 @@
-torchao==0.10.0
+torchao==0.17.0
 safetensors
 git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a
 #pip install git+https://github.com/huggingface/diffusers.git@refs/pull/13432/head

run_linux.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# Update-and-run script for Linux — thin bootstrap over the in-repo manager.
+#
+# Everything (venv via uv-managed Python, torch for your GPU, requirements,
+# portable Node.js and FFmpeg, dependency updates) is handled by
+# `python -m manager`; this script only makes sure uv + a Python interpreter
+# exist, then delegates. Works on desktop and headless boxes alike.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+# ── Banner ─────────────────────────────────────────────────────────
+echo ""
+printf '\033[36m'
+cat << 'BANNER'
+     _     ___   _____               _  _     _  _
+    / \   |_ _| |_   _|  ___    ___ | || | __(_)| |_
+   / _ \   | |    | |   / _ \  / _ \| || |/ /| || __|
+  / ___ \  | |    | |  | (_) || (_) | ||   < | || |_
+ /_/   \_\|___|   |_|   \___/  \___/|_||_|\_\|_| \__|
+BANNER
+printf '\033[0m'
+printf '\033[90m  Linux Setup & Launcher\033[0m\n'
+echo ""
+
+# ── 1. Ensure uv (prebuilt static binary, kept inside the repo) ─────
+export PATH="$SCRIPT_DIR/.uv:$PATH"
+export UV_PYTHON_INSTALL_DIR="$SCRIPT_DIR/.uv/python"
+if ! command -v uv >/dev/null 2>&1; then
+    echo "Downloading uv (package/python manager) into .uv/ ..."
+    curl -LsSf https://astral.sh/uv/install.sh | \
+        UV_INSTALL_DIR="$SCRIPT_DIR/.uv" UV_NO_MODIFY_PATH=1 sh
+fi
+
+# ── 2. Find a Python to run the manager (stdlib-only, needs >= 3.9) ─
+find_python() {
+    for cmd in python3 python; do
+        if command -v "$cmd" >/dev/null 2>&1; then
+            if "$cmd" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)' 2>/dev/null; then
+                echo "$cmd"
+                return 0
+            fi
+        fi
+    done
+    return 1
+}
+
+PYTHON="$(find_python || true)"
+if [[ -z "$PYTHON" ]]; then
+    echo "No system Python found — provisioning one with uv..."
+    uv python install 3.12
+    PYTHON="$(uv python find 3.12)"
+fi
+
+# ── 3. Sync the environment and start the UI ────────────────────────
+cd "$SCRIPT_DIR"
+"$PYTHON" -m manager update --auto
+exec "$PYTHON" -m manager launch

run_mac.zsh

@@ -1,5 +1,9 @@
 #!/usr/bin/env zsh
-# Update-and-run script for macOS — portable Python 3.12 + PyTorch
+# Update-and-run script for macOS — thin bootstrap over the in-repo manager.
+#
+# Everything (venv via uv-managed Python, torch, requirements, portable
+# Node.js and FFmpeg, dependency updates) is handled by `python -m manager`;
+# this script only makes sure uv + a Python interpreter exist, then delegates.
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
@@ -17,150 +21,37 @@ BANNER
 echo "\033[0m"
 echo "\033[90m  macOS Setup & Launcher\033[0m"
 echo ""
-VENV_DIR="$SCRIPT_DIR/.venv"
-PIP="$VENV_DIR/bin/pip"
-PYTHON="$VENV_DIR/bin/python3"
-PYTHON_VERSION="3.12.8"
-RELEASE_TAG="20241219"
 
-# --- Package versions (update these as needed) ---
+# ── 1. Ensure uv (prebuilt static binary, kept inside the repo) ─────
-NODE_VERSION="23.11.1"
+export PATH="$SCRIPT_DIR/.uv:$PATH"
-TORCH_VERSION="2.11.0"
+export UV_PYTHON_INSTALL_DIR="$SCRIPT_DIR/.uv/python"
-TORCHVISION_VERSION="0.26.0"
+if ! command -v uv >/dev/null 2>&1; then
-TORCHAUDIO_VERSION="2.11.0"
+    echo "Downloading uv (package/python manager) into .uv/ ..."
-
+    curl -LsSf https://astral.sh/uv/install.sh | \
-# Detect architecture
+        UV_INSTALL_DIR="$SCRIPT_DIR/.uv" UV_NO_MODIFY_PATH=1 sh
-ARCH="$(uname -m)"
-if [[ "$ARCH" == "arm64" ]]; then
-    PLATFORM="aarch64-apple-darwin"
-elif [[ "$ARCH" == "x86_64" ]]; then
-    PLATFORM="x86_64-apple-darwin"
-else
-    echo "Error: Unsupported architecture: $ARCH"
-    exit 1
 fi
 
-# ── 1. Download standalone Python if needed ─────────────────────────
+# ── 2. Find a Python to run the manager (stdlib-only, needs >= 3.9) ─
-PYTHON_DIR="$SCRIPT_DIR/.python"
+find_python() {
-PYTHON_BIN="$PYTHON_DIR/bin/python3"
+    for cmd in python3 python; do
-
+        if command -v "$cmd" >/dev/null 2>&1; then
-if [[ ! -x "$PYTHON_BIN" ]]; then
+            if "$cmd" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)' 2>/dev/null; then
-    TARBALL="cpython-${PYTHON_VERSION}+${RELEASE_TAG}-${PLATFORM}-install_only.tar.gz"
+                echo "$cmd"
-    URL="https://github.com/indygreg/python-build-standalone/releases/download/${RELEASE_TAG}/${TARBALL}"
+                return 0
-
+            fi
-    TMPDIR_DL="$(mktemp -d)"
+        fi
-    trap 'rm -rf "$TMPDIR_DL"' EXIT
+    done
-
+    return 1
-    echo "Downloading standalone Python ${PYTHON_VERSION} (${PLATFORM})..."
-    curl -fSL --progress-bar -o "$TMPDIR_DL/$TARBALL" "$URL"
-
-    echo "Extracting..."
-    tar -xzf "$TMPDIR_DL/$TARBALL" -C "$TMPDIR_DL"
-
-    # Move to permanent location (the archive extracts to a "python" folder)
-    rm -rf "$PYTHON_DIR"
-    mv "$TMPDIR_DL/python" "$PYTHON_DIR"
-
-    rm -rf "$TMPDIR_DL"
-    trap - EXIT
-
-    echo "Standalone Python installed to $PYTHON_DIR"
-fi
-
-# ── 2. Create venv if it doesn't exist ──────────────────────────────
-if [[ ! -d "$VENV_DIR" ]]; then
-    echo "Creating virtual environment at $VENV_DIR..."
-    "$PYTHON_BIN" -m venv "$VENV_DIR"
-    echo "Virtual environment created."
-fi
-
-# ── 3. Download / update portable Node.js ──────────────────────────
-NODE_DIR="$SCRIPT_DIR/.node"
-NODE_BIN="$NODE_DIR/bin/node"
-
-NEED_NODE=false
-if [[ ! -x "$NODE_BIN" ]]; then
-    NEED_NODE=true
-elif [[ "$("$NODE_BIN" --version 2>/dev/null)" != "v${NODE_VERSION}" ]]; then
-    echo "Node.js version mismatch (want v${NODE_VERSION}, have $("$NODE_BIN" --version))."
-    NEED_NODE=true
-fi
-
-if $NEED_NODE; then
-    if [[ "$ARCH" == "arm64" ]]; then
-        NODE_ARCH="arm64"
-    else
-        NODE_ARCH="x64"
-    fi
-
-    NODE_TARBALL="node-v${NODE_VERSION}-darwin-${NODE_ARCH}.tar.gz"
-    NODE_URL="https://nodejs.org/dist/v${NODE_VERSION}/${NODE_TARBALL}"
-
-    TMPDIR_DL="$(mktemp -d)"
-    trap 'rm -rf "$TMPDIR_DL"' EXIT
-
-    echo "Downloading Node.js v${NODE_VERSION} (darwin-${NODE_ARCH})..."
-    curl -fSL --progress-bar -o "$TMPDIR_DL/$NODE_TARBALL" "$NODE_URL"
-
-    echo "Extracting..."
-    tar -xzf "$TMPDIR_DL/$NODE_TARBALL" -C "$TMPDIR_DL"
-
-    rm -rf "$NODE_DIR"
-    mv "$TMPDIR_DL/node-v${NODE_VERSION}-darwin-${NODE_ARCH}" "$NODE_DIR"
-
-    rm -rf "$TMPDIR_DL"
-    trap - EXIT
-
-    echo "Node.js v${NODE_VERSION} installed to $NODE_DIR"
-else
-    echo "Node.js v${NODE_VERSION} is up to date."
-fi
-
-# ── 4. Install / update PyTorch packages ────────────────────────────
-# Helper: returns 0 if the package is installed at the exact version
-pkg_ok() {
-    local pkg="$1" want="$2"
-    local got
-    got="$("$PIP" show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')" || true
-    [[ "$got" == "$want" ]]
 }
 
-PKGS_TO_INSTALL=()
+PYTHON="$(find_python || true)"
-
+if [[ -z "$PYTHON" ]]; then
-pkg_ok "torch"       "$TORCH_VERSION"       || PKGS_TO_INSTALL+=("torch==$TORCH_VERSION")
+    echo "No system Python found — provisioning one with uv..."
-pkg_ok "torchvision" "$TORCHVISION_VERSION"  || PKGS_TO_INSTALL+=("torchvision==$TORCHVISION_VERSION")
+    uv python install 3.12
-pkg_ok "torchaudio"  "$TORCHAUDIO_VERSION"   || PKGS_TO_INSTALL+=("torchaudio==$TORCHAUDIO_VERSION")
+    PYTHON="$(uv python find 3.12)"
-
-if (( ${#PKGS_TO_INSTALL[@]} )); then
-    echo "Installing / updating: ${PKGS_TO_INSTALL[*]}"
-    "$PIP" install "${PKGS_TO_INSTALL[@]}"
-else
-    echo "PyTorch packages are up to date."
 fi
 
-# ── 5. Install / update requirements.txt ────────────────────────────
+# ── 3. Sync the environment and start the UI ────────────────────────
-REQUIREMENTS="$SCRIPT_DIR/requirements.txt"
+cd "$SCRIPT_DIR"
-REQ_HASH_FILE="$VENV_DIR/.requirements_hash"
+"$PYTHON" -m manager update --auto
-
+exec "$PYTHON" -m manager launch
-if [[ -f "$REQUIREMENTS" ]]; then
-    # Hash all requirements files (follows -r includes)
-    CURRENT_HASH="$(cat "$SCRIPT_DIR"/requirements*.txt 2>/dev/null | shasum -a 256 | awk '{print $1}')"
-    STORED_HASH=""
-    [[ -f "$REQ_HASH_FILE" ]] && STORED_HASH="$(cat "$REQ_HASH_FILE")"
-
-    if [[ "$CURRENT_HASH" != "$STORED_HASH" ]]; then
-        echo "Installing / updating requirements.txt..."
-        "$PIP" install -r "$REQUIREMENTS"
-        echo "$CURRENT_HASH" > "$REQ_HASH_FILE"
-    else
-        echo "Requirements are up to date."
-    fi
-fi
-
-# ── 6. Build and start the UI ───────────────────────────────────────
-export PATH="$NODE_DIR/bin:$VENV_DIR/bin:$PATH"
-
-echo ""
-echo "Starting UI..."
-cd "$SCRIPT_DIR/ui"
-npm run build_and_start

run_modal.py

@@ -55,7 +55,6 @@ image = (
         "toml", 
         "pydantic",
         "omegaconf",
-        "k-diffusion",
         "open_clip_torch",
         "timm",
         "prodigyopt",

run_windows.bat

@@ -0,0 +1,77 @@
+@echo off&&cd /d %~dp0
+REM Update-and-run script for Windows - thin bootstrap over the in-repo manager.
+REM
+REM Everything (venv via uv-managed Python, torch for your GPU, requirements,
+REM portable Node.js / FFmpeg / Git, dependency updates) is handled by
+REM `python -m manager`; this script only makes sure uv + a Python interpreter
+REM exist, then delegates.
+setlocal EnableDelayedExpansion
+Title AI Toolkit
+
+echo.
+echo      _     ___   _____               _  _     _  _
+echo     / \   ^|_ _^| ^|_   _^|  ___    ___ ^| ^|^| ^| __(_)^| ^|_
+echo    / _ \   ^| ^|    ^| ^|   / _ \  / _ \^| ^|^| ^|/ /^| ^|^| __^|
+echo   / ___ \  ^| ^|    ^| ^|  ^| (_) ^|^| (_) ^| ^|^|   ^< ^| ^|^| ^|_
+echo  /_/   \_\^|___^|   ^|_^|   \___/  \___/^|_^|^|_^|\_\^|_^| \__^|
+echo.
+echo   Windows Setup ^& Launcher
+echo.
+
+REM Clear env vars that let a stray conda/pyenv/system Python hijack things
+set PYTHONPATH=
+set PYTHONHOME=
+set PYTHONSTARTUP=
+set PYTHONUSERBASE=
+set PIP_CONFIG_FILE=
+set VIRTUAL_ENV=
+set CONDA_PREFIX=
+set CONDA_DEFAULT_ENV=
+set PYENV_ROOT=
+set PYENV_VERSION=
+
+REM ---- 1. Ensure uv (prebuilt static binary, kept inside the repo) ----
+set "PATH=%~dp0.uv;%PATH%"
+set "UV_PYTHON_INSTALL_DIR=%~dp0.uv\python"
+where uv.exe >nul 2>&1
+if errorlevel 1 (
+    echo Downloading uv ^(package/python manager^) into .uv\ ...
+    powershell -NoProfile -ExecutionPolicy ByPass -Command ^
+        "$env:UV_INSTALL_DIR = Join-Path '%~dp0' '.uv'; $env:UV_NO_MODIFY_PATH = '1'; irm https://astral.sh/uv/install.ps1 | iex"
+    where uv.exe >nul 2>&1
+    if errorlevel 1 (
+        echo ERROR: uv download failed. See https://docs.astral.sh/uv/
+        pause
+        exit /b 1
+    )
+)
+
+REM ---- 2. Find a Python to run the manager (stdlib-only, needs 3.9+) ----
+set "PY="
+for %%C in (python.exe py.exe) do (
+    if not defined PY (
+        %%C -c "import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)" >nul 2>&1
+        if not errorlevel 1 set "PY=%%C"
+    )
+)
+if not defined PY (
+    echo No system Python found - provisioning one with uv...
+    uv python install 3.12
+    for /f "delims=" %%P in ('uv python find 3.12') do set "PY=%%P"
+)
+if not defined PY (
+    echo ERROR: could not find or install a Python interpreter.
+    pause
+    exit /b 1
+)
+
+REM ---- 3. Sync the environment and start the UI ----
+"%PY%" -m manager update --auto
+if errorlevel 1 (
+    echo.
+    echo Setup failed - see output above.
+    pause
+    exit /b 1
+)
+"%PY%" -m manager launch
+pause