| 1 | """Install the pinned MSST-WebUI source and Echo 1.5 Bandit checkpoint.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import hashlib |
| 6 | import shutil |
| 7 | import subprocess |
| 8 | import tempfile |
| 9 | import urllib.request |
| 10 | from pathlib import Path |
| 11 | |
| 12 | |
| 13 | REPO_ROOT = Path(__file__).resolve().parents[1] |
| 14 | MSST_REPOSITORY = "https://github.com/SUC-DriverOld/MSST-WebUI.git" |
| 15 | MSST_COMMIT = "43e30b860c611b516ed9b67c75a56792a67ec902" |
| 16 | MSST_DIR = REPO_ROOT / "third_party" / "MSST-WebUI" |
| 17 | MODEL_NAME = "model_bandit_plus_dnr_sdr_11.47.chpt" |
| 18 | MODEL_URL = ( |
| 19 | "https://huggingface.co/Sucial/MSST-WebUI/resolve/main/" |
| 20 | f"All_Models/multi_stem_models/{MODEL_NAME}" |
| 21 | ) |
| 22 | MODEL_SHA256 = "c48284779f7d1258a6527d3aaa18a532d45c1f506e2dcc25d5ab179a8c5e2573" |
| 23 | MODEL_PATH = REPO_ROOT / "checkpoints" / "msst" / MODEL_NAME |
| 24 | CONFIG_PATH = MSST_DIR / "configs_backup" / "multi_stem_models" / f"{MODEL_NAME}.yaml" |
| 25 | CONFIG_SHA256 = "4d3bf5b9fb9d0480bf9cb64eaed23edb37815aaaf905d74ca6393c13a32ce58e" |
| 26 | |
| 27 | |
| 28 | def _run(*args: str, cwd: Path | None = None) -> str: |
| 29 | result = subprocess.run( |
| 30 | list(args), |
| 31 | cwd=cwd, |
| 32 | check=True, |
| 33 | capture_output=True, |
| 34 | text=True, |
| 35 | ) |
| 36 | return result.stdout.strip() |
| 37 | |
| 38 | |
| 39 | def _sha256(path: Path) -> str: |
| 40 | digest = hashlib.sha256() |
| 41 | with path.open("rb") as handle: |
| 42 | for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| 43 | digest.update(chunk) |
| 44 | return digest.hexdigest() |
| 45 | |
| 46 | |
| 47 | def install_source() -> None: |
| 48 | if MSST_DIR.exists(): |
| 49 | if not (MSST_DIR / ".git").is_dir(): |
| 50 | raise RuntimeError(f"existing MSST path is not a Git checkout: {MSST_DIR}") |
| 51 | current_commit = _run("git", "rev-parse", "HEAD", cwd=MSST_DIR) |
| 52 | if current_commit != MSST_COMMIT: |
| 53 | raise RuntimeError( |
| 54 | f"MSST checkout is at {current_commit}; expected {MSST_COMMIT}. " |
| 55 | "Move it aside before running setup again." |
| 56 | ) |
| 57 | if _run("git", "status", "--porcelain", cwd=MSST_DIR): |
| 58 | raise RuntimeError(f"MSST checkout contains local changes: {MSST_DIR}") |
| 59 | return |
| 60 | |
| 61 | MSST_DIR.parent.mkdir(parents=True, exist_ok=True) |
| 62 | _run( |
| 63 | "git", |
| 64 | "clone", |
| 65 | "--filter=blob:none", |
| 66 | "--no-checkout", |
| 67 | MSST_REPOSITORY, |
| 68 | str(MSST_DIR), |
| 69 | ) |
| 70 | _run("git", "checkout", "--detach", MSST_COMMIT, cwd=MSST_DIR) |
| 71 | |
| 72 | |
| 73 | def install_model() -> None: |
| 74 | MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) |
| 75 | if MODEL_PATH.is_file(): |
| 76 | actual = _sha256(MODEL_PATH) |
| 77 | if actual != MODEL_SHA256: |
| 78 | raise RuntimeError( |
| 79 | f"existing model checksum mismatch: expected {MODEL_SHA256}, got {actual}" |
| 80 | ) |
| 81 | return |
| 82 | |
| 83 | temporary_path: Path | None = None |
| 84 | try: |
| 85 | with tempfile.NamedTemporaryFile( |
| 86 | dir=MODEL_PATH.parent, |
| 87 | prefix=f".{MODEL_NAME}.", |
| 88 | suffix=".download", |
| 89 | delete=False, |
| 90 | ) as temporary: |
| 91 | temporary_path = Path(temporary.name) |
| 92 | request = urllib.request.Request( |
| 93 | MODEL_URL, |
| 94 | headers={"User-Agent": "JoyAI-Echo15-MSST-Setup/1.0"}, |
| 95 | ) |
| 96 | with urllib.request.urlopen(request) as response: |
| 97 | shutil.copyfileobj(response, temporary) |
| 98 | actual = _sha256(temporary_path) |
| 99 | if actual != MODEL_SHA256: |
| 100 | raise RuntimeError( |
| 101 | f"downloaded model checksum mismatch: expected {MODEL_SHA256}, got {actual}" |
| 102 | ) |
| 103 | temporary_path.replace(MODEL_PATH) |
| 104 | temporary_path = None |
| 105 | finally: |
| 106 | if temporary_path is not None: |
| 107 | temporary_path.unlink(missing_ok=True) |
| 108 | |
| 109 | |
| 110 | def main() -> None: |
| 111 | install_source() |
| 112 | install_model() |
| 113 | if not CONFIG_PATH.is_file(): |
| 114 | raise FileNotFoundError(f"MSST Bandit config not found: {CONFIG_PATH}") |
| 115 | config_digest = _sha256(CONFIG_PATH) |
| 116 | if config_digest != CONFIG_SHA256: |
| 117 | raise RuntimeError( |
| 118 | f"MSST config checksum mismatch: expected {CONFIG_SHA256}, got {config_digest}" |
| 119 | ) |
| 120 | print(f"MSST source: {MSST_DIR}") |
| 121 | print(f"MSST commit: {MSST_COMMIT}") |
| 122 | print(f"MSST model: {MODEL_PATH}") |
| 123 | print(f"MSST config: {CONFIG_PATH}") |
| 124 | |
| 125 | |
| 126 | if __name__ == "__main__": |
| 127 | main() |
| 128 |