use local model cache

examples/01_chat.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 os.environ.setdefault("DEV", "CPU")
-os.environ.setdefault("XDG_CACHE_HOME", str(Path(__file__).resolve().parents[1] / ".cache"))
+os.environ["XDG_CACHE_HOME"] = str(Path(__file__).resolve().parents[1] / ".cache")
 
 from tinygrad.llm.cli import main
 if __name__ == "__main__":

examples/04_classify.py

@@ -4,7 +4,7 @@
 
 ROOT = Path(__file__).resolve().parents[1]
 os.environ.setdefault('DEV', 'CPU')
-os.environ.setdefault('XDG_CACHE_HOME', str(ROOT / '.cache'))
+os.environ['XDG_CACHE_HOME'] = str(ROOT / '.cache')
 
 from PIL import Image
 from torchvision.models import ResNet18_Weights

examples/vision.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 
 os.environ.setdefault("DEV", "CPU")
-os.environ.setdefault("XDG_CACHE_HOME", str(Path(__file__).resolve().parents[1] / ".cache"))
+os.environ["XDG_CACHE_HOME"] = str(Path(__file__).resolve().parents[1] / ".cache")
 
 import numpy as np
 import torch

setup.sh

@@ -2,12 +2,18 @@
 set -euo pipefail
 cd "$(dirname "$0")"
 export PATH="$HOME/.local/bin:/opt/homebrew/bin:$PATH"
-export UV_CACHE_DIR="${UV_CACHE_DIR:-$PWD/.cache/uv}"
+export UV_CACHE_DIR="$PWD/.cache/uv"
+export TORCH_HOME="$PWD/.cache/torch"
 export UV_HTTP_TIMEOUT=200 UV_HTTP_RETRIES=5
 mkdir -p .cache
 
 as_root=()
 if [ "$EUID" -ne 0 ]; then as_root=(sudo); fi
+for directory in .cache .venv models frames; do
+  if [ -d "$directory" ] && [ -n "$(find "$directory" ! -user "$(id -un)" -print -quit)" ]; then
+    "${as_root[@]}" chown -R "$(id -u):$(id -g)" "$directory"
+  fi
+done
 download() { curl --retry 5 --retry-delay 5 --retry-all-errors -fLsS "$1" -o "$2"; }
 
 case "$(uname -s):$(uname -m)" in
@@ -79,7 +85,9 @@ if ! command -v uv >/dev/null; then
   download https://astral.sh/uv/install.sh .cache/install-uv.sh
   UV_NO_MODIFY_PATH=1 sh .cache/install-uv.sh
 fi
+unset VIRTUAL_ENV
 uv sync --locked --python 3.12
 .venv/bin/python tools/setup.py
+echo 'Preparing models...'
 .venv/bin/python tools/export.py
 echo 'Ready. Activate with: source .venv/bin/activate'

tools/export.py

@@ -13,3 +13,4 @@
   model = resnet18(weights=ResNet18_Weights.DEFAULT).eval()
   torch.onnx.export(model, torch.zeros(1, 3, 224, 224), folder / 'resnet18.onnx',
                     input_names=['image'], output_names=['scores'], opset_version=17, dynamo=False)
+print('Models ready.', flush=True)