212 lines
7.1 KiB
Python
212 lines
7.1 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
|
|
MODELS_BASE_URLS = (
|
|
"https://git.konn3kt.com/teal/IQModels/raw/branch/main",
|
|
"https://gitlvb.teallvbs.xyz/teal/IQModels/raw/branch/main",
|
|
)
|
|
CHUNK = 4 * 1024 * 1024
|
|
HTTP_TIMEOUT_S = 60.0
|
|
STREAM_RETRIES = 6
|
|
|
|
|
|
def _requests_auth():
|
|
import importlib
|
|
for mod in ("iqpilot_private.models.git_auth", "iqpilot.models_private_src.git_auth",
|
|
"iqpilot.selfdrive.iqmodeld.models.git_auth"):
|
|
try:
|
|
return importlib.import_module(mod).get_requests_auth()
|
|
except Exception:
|
|
continue
|
|
return None
|
|
|
|
|
|
def _hf():
|
|
import importlib
|
|
for mod in ("iqpilot_private.models.git_auth", "iqpilot.selfdrive.iqmodeld.models.git_auth"):
|
|
try:
|
|
m = importlib.import_module(mod)
|
|
return m.get_hf_headers(), m.hf_resolve_url
|
|
except Exception:
|
|
continue
|
|
return None, None
|
|
|
|
|
|
def download_hf_file(hf_path: str, dst: str, sha256: str, size: int, progress_cb=None) -> str:
|
|
import requests
|
|
headers, resolve = _hf()
|
|
if resolve is None:
|
|
raise RuntimeError("no HF credentials available")
|
|
url = resolve(hf_path)
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
tmp = dst + ".hfpart"
|
|
last_error: Exception | None = None
|
|
for _attempt in range(STREAM_RETRIES):
|
|
try:
|
|
have = os.path.getsize(tmp) if os.path.isfile(tmp) else 0
|
|
if size and have > size:
|
|
os.remove(tmp)
|
|
have = 0
|
|
if not size or have < size:
|
|
req_headers = dict(headers)
|
|
if have:
|
|
req_headers["Range"] = f"bytes={have}-"
|
|
with requests.get(url, headers=req_headers, stream=True, timeout=HTTP_TIMEOUT_S, allow_redirects=True) as r:
|
|
r.raise_for_status()
|
|
if have and r.status_code != 206:
|
|
have = 0
|
|
with open(tmp, "ab" if have else "wb") as f:
|
|
got = have
|
|
for chunk in r.iter_content(CHUNK):
|
|
f.write(chunk)
|
|
got += len(chunk)
|
|
if progress_cb is not None and size:
|
|
progress_cb(min(1.0, got / size))
|
|
digest = hashlib.sha256()
|
|
with open(tmp, "rb") as f:
|
|
for chunk in iter(lambda: f.read(CHUNK), b""):
|
|
digest.update(chunk)
|
|
if size and os.path.getsize(tmp) != size:
|
|
raise RuntimeError(f"size mismatch: {os.path.getsize(tmp)}/{size} bytes")
|
|
if sha256 and digest.hexdigest() != sha256:
|
|
os.remove(tmp)
|
|
raise RuntimeError("sha256 mismatch")
|
|
os.replace(tmp, dst)
|
|
return dst
|
|
except Exception as e:
|
|
last_error = e
|
|
raise RuntimeError(f"HF download failed: {last_error}")
|
|
|
|
|
|
def _lfs_endpoint(base_url: str) -> str:
|
|
return base_url.split("/raw/", 1)[0] + ".git/info/lfs"
|
|
|
|
|
|
def _resolve_oid(session, base_url: str, oid: str, size: int, auth):
|
|
import requests
|
|
batch = session.post(f"{_lfs_endpoint(base_url)}/objects/batch",
|
|
data=json.dumps({"operation": "download", "transfers": ["basic"],
|
|
"objects": [{"oid": oid, "size": size}]}),
|
|
headers={"Content-Type": "application/vnd.git-lfs+json",
|
|
"Accept": "application/vnd.git-lfs+json"},
|
|
auth=auth, timeout=HTTP_TIMEOUT_S)
|
|
batch.raise_for_status()
|
|
entry = batch.json()["objects"][0]
|
|
if "actions" not in entry:
|
|
raise requests.RequestException(f"LFS object unavailable: {entry.get('error', oid)}")
|
|
action = entry["actions"]["download"]
|
|
return action["href"], action.get("header", {})
|
|
|
|
|
|
def _part_path(dst: str, oid: str) -> str:
|
|
return os.path.join(dst + ".parts", oid)
|
|
|
|
|
|
def _part_complete(path: str, oid: str, size: int) -> bool:
|
|
if not os.path.isfile(path) or os.path.getsize(path) != size:
|
|
return False
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(CHUNK), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest() == oid
|
|
|
|
|
|
def _fetch_part(session, base_url: str, obj: dict, path: str, auth, progress) -> None:
|
|
size = int(obj["size"])
|
|
have = os.path.getsize(path) if os.path.isfile(path) else 0
|
|
if have > size:
|
|
os.remove(path)
|
|
have = 0
|
|
href, headers = _resolve_oid(session, base_url, obj["oid"], size, auth)
|
|
obj_auth = None if headers.get("Authorization") else auth
|
|
# LFS parts are content-addressed (oid == sha256), so a half-written part can be resumed with a
|
|
# Range request and verified afterwards instead of being thrown away on every restart.
|
|
if have:
|
|
headers = {**headers, "Range": f"bytes={have}-"}
|
|
with session.get(href, headers=headers, stream=True, timeout=HTTP_TIMEOUT_S, auth=obj_auth) as r:
|
|
r.raise_for_status()
|
|
if have and r.status_code != 206:
|
|
have = 0
|
|
with open(path, "ab" if have else "wb") as f:
|
|
for chunk in r.iter_content(CHUNK):
|
|
f.write(chunk)
|
|
progress(len(chunk))
|
|
|
|
|
|
def download_lfs_bundle(objects: list, dst: str, sha256: str, size: int, progress_cb=None) -> str:
|
|
import requests
|
|
auth = _requests_auth()
|
|
session = requests.Session()
|
|
os.makedirs(dst + ".parts", exist_ok=True)
|
|
total = int(size) or sum(int(o["size"]) for o in objects)
|
|
done_bytes = sum(int(o["size"]) for o in objects if _part_complete(_part_path(dst, o["oid"]), o["oid"], int(o["size"])))
|
|
got = [done_bytes]
|
|
|
|
def progress(n: int) -> None:
|
|
got[0] += n
|
|
if progress_cb is not None and total:
|
|
progress_cb(min(1.0, got[0] / total))
|
|
|
|
last_error: Exception | None = None
|
|
for base_url in MODELS_BASE_URLS:
|
|
for _attempt in range(STREAM_RETRIES):
|
|
try:
|
|
for obj in objects:
|
|
path = _part_path(dst, obj["oid"])
|
|
if _part_complete(path, obj["oid"], int(obj["size"])):
|
|
continue
|
|
got[0] = done_bytes
|
|
_fetch_part(session, base_url, obj, path, auth, progress)
|
|
if not _part_complete(path, obj["oid"], int(obj["size"])):
|
|
if os.path.getsize(path) >= int(obj["size"]):
|
|
os.remove(path)
|
|
raise RuntimeError(f"part {obj['oid'][:12]} incomplete or failed verification")
|
|
done_bytes += int(obj["size"])
|
|
got[0] = done_bytes
|
|
break
|
|
except Exception as e:
|
|
last_error = e
|
|
else:
|
|
continue
|
|
break
|
|
else:
|
|
raise RuntimeError(f"model bundle download failed: {last_error}")
|
|
|
|
tmp = dst + ".part"
|
|
digest = hashlib.sha256()
|
|
with open(tmp, "wb") as out:
|
|
for obj in objects:
|
|
with open(_part_path(dst, obj["oid"]), "rb") as f:
|
|
for chunk in iter(lambda: f.read(CHUNK), b""):
|
|
out.write(chunk)
|
|
digest.update(chunk)
|
|
if total and os.path.getsize(tmp) != total:
|
|
os.remove(tmp)
|
|
raise RuntimeError(f"size mismatch: {os.path.getsize(tmp) if os.path.exists(tmp) else 0}/{total} bytes")
|
|
if sha256 and digest.hexdigest() != sha256:
|
|
os.remove(tmp)
|
|
for obj in objects:
|
|
try:
|
|
os.remove(_part_path(dst, obj["oid"]))
|
|
except OSError:
|
|
pass
|
|
raise RuntimeError("sha256 mismatch")
|
|
os.replace(tmp, dst)
|
|
for obj in objects:
|
|
try:
|
|
os.remove(_part_path(dst, obj["oid"]))
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.rmdir(dst + ".parts")
|
|
except OSError:
|
|
pass
|
|
return dst
|