Initial commit
This commit is contained in:
+314
@@ -0,0 +1,314 @@
|
||||
"""First-run launcher.
|
||||
|
||||
Bootstraps a clean Python 3.11 environment via `uv`, regardless of the system
|
||||
Python the user invoked us with. Keeps the user on a single supported runtime
|
||||
while the AI ecosystem stabilizes around newer Python versions.
|
||||
|
||||
uv install strategy (in order):
|
||||
1. Direct binary download from GitHub releases (no shell, no admin).
|
||||
2. PowerShell / shell installer.
|
||||
3. pip fallback, invoked as `python -m uv`.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).parent.resolve()
|
||||
VENV_DIR = ROOT / "venv"
|
||||
MARKER = VENV_DIR / ".kawai_ready"
|
||||
HARDWARE_CACHE = ROOT / "config.local.json"
|
||||
UV_CACHE_DIR = ROOT / ".tools"
|
||||
PYTHON_TARGET = "3.11"
|
||||
|
||||
|
||||
def venv_python() -> Path:
|
||||
if os.name == "nt":
|
||||
return VENV_DIR / "Scripts" / "python.exe"
|
||||
return VENV_DIR / "bin" / "python"
|
||||
|
||||
|
||||
# --- uv discovery / install ------------------------------------------------
|
||||
|
||||
def _local_uv_path() -> Path:
|
||||
return UV_CACHE_DIR / ("uv.exe" if os.name == "nt" else "uv")
|
||||
|
||||
|
||||
def _uv_argv() -> list[str] | None:
|
||||
"""Return argv prefix to invoke uv. None if not available."""
|
||||
local = _local_uv_path()
|
||||
if local.exists():
|
||||
return [str(local)]
|
||||
found = shutil.which("uv")
|
||||
if found:
|
||||
return [found]
|
||||
# Installed via pip into current Python (works as module).
|
||||
try:
|
||||
subprocess.check_output(
|
||||
[sys.executable, "-m", "uv", "--version"],
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=10,
|
||||
)
|
||||
return [sys.executable, "-m", "uv"]
|
||||
except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
|
||||
|
||||
def _uv_release_asset() -> str | None:
|
||||
"""Pick the GitHub release asset name for this OS+arch."""
|
||||
machine = platform.machine().lower()
|
||||
arch_win = "x86_64" if machine in ("amd64", "x86_64") else "aarch64" if "arm" in machine else None
|
||||
if os.name == "nt" and arch_win:
|
||||
return f"uv-{arch_win}-pc-windows-msvc.zip"
|
||||
if sys.platform == "darwin":
|
||||
arch = "aarch64" if "arm" in machine else "x86_64"
|
||||
return f"uv-{arch}-apple-darwin.tar.gz"
|
||||
if sys.platform.startswith("linux"):
|
||||
arch = "aarch64" if "aarch64" in machine or "arm64" in machine else "x86_64"
|
||||
return f"uv-{arch}-unknown-linux-gnu.tar.gz"
|
||||
return None
|
||||
|
||||
|
||||
def _download_uv() -> bool:
|
||||
"""Download uv binary directly from GitHub releases. Most reliable path."""
|
||||
asset = _uv_release_asset()
|
||||
if asset is None:
|
||||
return False
|
||||
url = f"https://github.com/astral-sh/uv/releases/latest/download/{asset}"
|
||||
UV_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
print(f"[kawai] Downloading uv from {url}")
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=60) as resp:
|
||||
data = resp.read()
|
||||
except Exception as e:
|
||||
print(f"[kawai] download failed: {e}")
|
||||
return False
|
||||
|
||||
try:
|
||||
if asset.endswith(".zip"):
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as z:
|
||||
for name in z.namelist():
|
||||
if name.endswith("uv.exe") or name.endswith("/uv"):
|
||||
target = _local_uv_path()
|
||||
target.write_bytes(z.read(name))
|
||||
if os.name != "nt":
|
||||
target.chmod(0o755)
|
||||
return target.exists()
|
||||
else:
|
||||
import tarfile
|
||||
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as t:
|
||||
for member in t.getmembers():
|
||||
if member.name.endswith("/uv") or member.name == "uv":
|
||||
f = t.extractfile(member)
|
||||
if f is None:
|
||||
continue
|
||||
target = _local_uv_path()
|
||||
target.write_bytes(f.read())
|
||||
target.chmod(0o755)
|
||||
return target.exists()
|
||||
except Exception as e:
|
||||
print(f"[kawai] extract failed: {e}")
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def _install_uv_via_shell() -> bool:
|
||||
"""Use astral.sh installer scripts. Often blocked on locked-down systems."""
|
||||
UV_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
env = os.environ.copy()
|
||||
env["UV_INSTALL_DIR"] = str(UV_CACHE_DIR)
|
||||
env["UV_NO_MODIFY_PATH"] = "1"
|
||||
try:
|
||||
if os.name == "nt":
|
||||
subprocess.check_call(
|
||||
[
|
||||
"powershell",
|
||||
"-NoProfile",
|
||||
"-ExecutionPolicy", "Bypass",
|
||||
"-Command",
|
||||
"irm https://astral.sh/uv/install.ps1 | iex",
|
||||
],
|
||||
env=env,
|
||||
)
|
||||
else:
|
||||
subprocess.check_call(
|
||||
["bash", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"],
|
||||
env=env,
|
||||
)
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return False
|
||||
return _local_uv_path().exists()
|
||||
|
||||
|
||||
def _install_uv_via_pip() -> bool:
|
||||
print("[kawai] Installing uv via pip...")
|
||||
try:
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", "--user", "--upgrade", "uv"])
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
# Verify it's invokable as a module.
|
||||
try:
|
||||
subprocess.check_output(
|
||||
[sys.executable, "-m", "uv", "--version"],
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=10,
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _ensure_uv() -> list[str]:
|
||||
argv = _uv_argv()
|
||||
if argv:
|
||||
return argv
|
||||
|
||||
print("[kawai] Installing uv...")
|
||||
if _download_uv():
|
||||
return _uv_argv() or []
|
||||
|
||||
if _install_uv_via_shell():
|
||||
argv = _uv_argv()
|
||||
if argv:
|
||||
return argv
|
||||
|
||||
if _install_uv_via_pip():
|
||||
argv = _uv_argv()
|
||||
if argv:
|
||||
return argv
|
||||
|
||||
raise RuntimeError(
|
||||
"Failed to install uv. Install manually from https://astral.sh/uv "
|
||||
"and rerun this launcher."
|
||||
)
|
||||
|
||||
|
||||
# --- venv + deps -----------------------------------------------------------
|
||||
|
||||
def _create_venv(uv: list[str]) -> None:
|
||||
if venv_python().exists():
|
||||
return
|
||||
print(f"[kawai] Creating venv with Python {PYTHON_TARGET} (uv will download it if needed)...")
|
||||
subprocess.check_call([*uv, "venv", str(VENV_DIR), "--python", PYTHON_TARGET])
|
||||
|
||||
|
||||
def _uv_pip(uv: list[str], args: list[str]) -> None:
|
||||
cmd = [*uv, "pip", "install", "--python", str(venv_python()), *args]
|
||||
print(f"[kawai] uv pip install {' '.join(args)}")
|
||||
subprocess.check_call(cmd)
|
||||
|
||||
|
||||
def detect_and_install(
|
||||
uv: list[str],
|
||||
force_backend: str | None = None,
|
||||
force_vendor: str | None = None,
|
||||
) -> dict:
|
||||
sys.path.insert(0, str(ROOT))
|
||||
from backends import hardware
|
||||
|
||||
info = hardware.detect(force_backend=force_backend, force_vendor=force_vendor)
|
||||
forced_note = " (forced)" if force_backend and force_backend != "auto" else ""
|
||||
print(
|
||||
f"[kawai] Backend: {info.backend}{forced_note} | "
|
||||
f"{info.vendor} / {info.device_name} / {info.vram_gb:.1f} GB / tier={info.tier}"
|
||||
)
|
||||
|
||||
_uv_pip(uv, hardware.torch_install_args(info))
|
||||
_uv_pip(uv, ["-r", str(ROOT / "requirements.txt")])
|
||||
|
||||
payload = {
|
||||
"vendor": info.vendor,
|
||||
"backend": info.backend,
|
||||
"device_name": info.device_name,
|
||||
"vram_gb": info.vram_gb,
|
||||
"tier": info.tier,
|
||||
"forced": bool(force_backend and force_backend != "auto"),
|
||||
}
|
||||
HARDWARE_CACHE.write_text(json.dumps(payload, indent=2))
|
||||
MARKER.write_text("ok")
|
||||
return payload
|
||||
|
||||
|
||||
def already_in_venv() -> bool:
|
||||
try:
|
||||
return Path(sys.executable).resolve() == venv_python().resolve()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def relaunch_in_venv(forwarded_args: list[str]) -> None:
|
||||
"""Re-exec the launcher inside the venv. Use subprocess on Windows because
|
||||
os.execv mangles argv with spaces in paths."""
|
||||
print("[kawai] Relaunching inside venv...")
|
||||
py = str(venv_python())
|
||||
script = str(ROOT / "launcher.py")
|
||||
argv = [py, script, *forwarded_args]
|
||||
if os.name == "nt":
|
||||
result = subprocess.run(argv)
|
||||
sys.exit(result.returncode)
|
||||
else:
|
||||
os.execv(py, argv)
|
||||
|
||||
|
||||
def _build_arg_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(
|
||||
prog="kawai",
|
||||
description="Local AI image/video generator. Auto-detects GPU; pass --backend to override.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--backend",
|
||||
choices=["auto", "cuda", "rocm", "directml", "mps", "cpu"],
|
||||
default="auto",
|
||||
help=(
|
||||
"Force torch backend. cuda=NVIDIA, rocm=AMD on Linux, directml=AMD/Intel on Windows, "
|
||||
"mps=Apple Silicon, cpu=fallback. Default: auto-detect."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--vendor",
|
||||
choices=["nvidia", "amd", "intel", "apple", "cpu"],
|
||||
default=None,
|
||||
help="Override detected vendor (rarely needed; useful when pairing --backend directml with intel).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--reinstall",
|
||||
action="store_true",
|
||||
help="Force re-detect and reinstall torch (clears the install marker).",
|
||||
)
|
||||
return p
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_arg_parser().parse_args()
|
||||
forced = args.backend if args.backend != "auto" else None
|
||||
|
||||
if args.reinstall and MARKER.exists():
|
||||
print("[kawai] --reinstall: clearing install marker")
|
||||
MARKER.unlink()
|
||||
|
||||
if already_in_venv():
|
||||
if not MARKER.exists():
|
||||
uv = _ensure_uv()
|
||||
detect_and_install(uv, force_backend=forced, force_vendor=args.vendor)
|
||||
from app import run
|
||||
run()
|
||||
return
|
||||
|
||||
uv = _ensure_uv()
|
||||
_create_venv(uv)
|
||||
if not MARKER.exists():
|
||||
detect_and_install(uv, force_backend=forced, force_vendor=args.vendor)
|
||||
relaunch_in_venv(sys.argv[1:])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user