95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
|
|
MANIFEST_PATH = os.path.join(ROOT, 'ai/model.json')
|
|
|
|
with open(MANIFEST_PATH, 'r', encoding='utf-8') as f:
|
|
manifest = json.load(f)
|
|
|
|
dest_path = os.path.join(ROOT, 'ai', manifest['file'])
|
|
part_path = dest_path + '.part'
|
|
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
|
|
|
expected_bytes = manifest['bytes']
|
|
expected_sha256 = manifest['sha256'].lower()
|
|
|
|
if os.path.exists(dest_path):
|
|
actual_bytes = os.path.getsize(dest_path)
|
|
if actual_bytes == expected_bytes:
|
|
print(f"[INFO] Destination file already exists ({actual_bytes} bytes). Verifying SHA256...")
|
|
h = hashlib.sha256()
|
|
with open(dest_path, 'rb') as f:
|
|
while chunk := f.read(1024 * 1024):
|
|
h.update(chunk)
|
|
if h.hexdigest().lower() == expected_sha256:
|
|
print("[SUCCESS] Model already downloaded and verified!")
|
|
sys.exit(0)
|
|
else:
|
|
print("[WARN] Existing file has checksum mismatch, re-downloading...")
|
|
os.remove(dest_path)
|
|
|
|
url = f"https://huggingface.co/{manifest['source']}/resolve/{manifest['revision']}/{manifest['source_file']}?download=true"
|
|
print(f"Downloading model from:\n{url}")
|
|
print(f"Destination: {dest_path}")
|
|
print(f"Expected size: {expected_bytes} bytes (~{expected_bytes / (1024**3):.2f} GB)")
|
|
|
|
headers = {'User-Agent': 'Mozilla/5.0'}
|
|
existing_bytes = 0
|
|
if os.path.exists(part_path):
|
|
existing_bytes = os.path.getsize(part_path)
|
|
if existing_bytes < expected_bytes:
|
|
print(f"Resuming download from byte {existing_bytes}...")
|
|
headers['Range'] = f"bytes={existing_bytes}-"
|
|
else:
|
|
existing_bytes = 0
|
|
os.remove(part_path)
|
|
|
|
req = urllib.request.Request(url, headers=headers)
|
|
try:
|
|
response = urllib.request.urlopen(req, timeout=30)
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to open URL: {e}")
|
|
sys.exit(1)
|
|
|
|
mode = 'ab' if existing_bytes > 0 else 'wb'
|
|
downloaded = existing_bytes
|
|
start_time = time.time()
|
|
last_report_time = start_time
|
|
last_report_bytes = downloaded
|
|
|
|
with open(part_path, mode) as out:
|
|
while True:
|
|
chunk = response.read(1024 * 512) # 512 KB
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
downloaded += len(chunk)
|
|
now = time.time()
|
|
if now - last_report_time >= 5.0 or downloaded == expected_bytes:
|
|
speed_mb = (downloaded - last_report_bytes) / (now - last_report_time) / (1024 * 1024)
|
|
pct = (downloaded / expected_bytes) * 100.0
|
|
print(f"Progress: {downloaded / (1024**3):.2f} / {expected_bytes / (1024**3):.2f} GB ({pct:.1f}%) | {speed_mb:.2f} MB/s")
|
|
last_report_time = now
|
|
last_report_bytes = downloaded
|
|
|
|
print("\nDownload finished! Verifying SHA256 checksum...")
|
|
h = hashlib.sha256()
|
|
with open(part_path, 'rb') as f:
|
|
while chunk := f.read(1024 * 1024):
|
|
h.update(chunk)
|
|
|
|
actual_sha256 = h.hexdigest().lower()
|
|
if actual_sha256 != expected_sha256:
|
|
print(f"[ERROR] SHA256 mismatch!\nExpected: {expected_sha256}\nGot: {actual_sha256}")
|
|
sys.exit(1)
|
|
|
|
if os.path.exists(dest_path):
|
|
os.remove(dest_path)
|
|
os.rename(part_path, dest_path)
|
|
print("[SUCCESS] Model successfully downloaded and verified!")
|