Files
InkFlow/backend/scripts/setup_models.py

88 lines
2.5 KiB
Python

#!/usr/bin/env python
"""Verifie l'environnement InkFlow et pre-telecharge les modeles MLX.
Usage :
python scripts/setup_models.py # tout verifier + telecharger
python scripts/setup_models.py --check # verifier sans telecharger
Pre-requis systeme : Apple Silicon, Python >= 3.11, ffmpeg (brew install ffmpeg).
"""
from __future__ import annotations
import argparse
import platform
import shutil
import sys
# Permet de lancer le script directement depuis backend/.
sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve().parents[1]))
from inkflow.config import ( # noqa: E402
GEMMA_MODEL,
KOKORO_MODEL,
QWEN3_TTS_MODEL,
ensure_dirs,
)
def check_env() -> bool:
ok = True
print(f"• Plateforme : {platform.platform()} ({platform.machine()})")
if platform.machine() != "arm64":
print(" ! Attendu arm64 (Apple Silicon) — MLX ne sera pas optimal.")
print(f"• Python : {sys.version.split()[0]}")
if sys.version_info < (3, 11):
print(" ! Python >= 3.11 requis."); ok = False
for mod in ("mlx", "mlx_lm", "mlx_audio", "ebooklib", "bs4",
"soundfile", "mutagen", "fastapi"):
try:
__import__(mod)
print(f"• import {mod:12s}: OK")
except Exception as exc: # noqa: BLE001
print(f"• import {mod:12s}: ECHEC ({exc})"); ok = False
ff = shutil.which("ffmpeg")
print(f"• ffmpeg : {ff or 'INTROUVABLE — brew install ffmpeg'}")
ok = ok and bool(ff)
return ok
def download_lm(model_id: str) -> None:
from mlx_lm import load
print(f" -> LM {model_id}")
load(model_id)
def download_tts(model_id: str) -> None:
from mlx_audio.tts.utils import load_model
print(f" -> TTS {model_id}")
load_model(model_id)
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--check", action="store_true", help="verifier sans telecharger")
args = ap.parse_args()
ensure_dirs()
print("== Verification de l'environnement ==")
env_ok = check_env()
if args.check:
return 0 if env_ok else 1
if not env_ok:
print("\nEnvironnement incomplet — corrige les points ci-dessus avant de continuer.")
return 1
print("\n== Telechargement des modeles (peut etre long la 1re fois) ==")
download_lm(GEMMA_MODEL)
download_tts(KOKORO_MODEL)
download_tts(QWEN3_TTS_MODEL)
print("\nTout est pret.")
return 0
if __name__ == "__main__":
raise SystemExit(main())