File size: 1,964 Bytes
af84b0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """Fetch the official Medium checkpoint through a mirror and verify its SHA-256."""
from concurrent.futures import ThreadPoolExecutor
import hashlib
from pathlib import Path
import subprocess
url = ("https://hf-mirror.com/wondervictor/YOLO-World/resolve/main/"
"yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.pth?download=true")
size, block = 367602867, 4 * 1024 * 1024
expected = "c6237d5b8f6cdc02c7e1a49ebca5ca6af2acb40dcc35908f3ac9eca5bee1120f"
root = Path(__file__).resolve().parent
directory = root / "weights" / "parts"
directory.mkdir(exist_ok=True)
def fetch(start):
end = min(start + block, size) - 1
path = directory / str(start)
for attempt in range(5):
have = path.stat().st_size if path.exists() else 0
if have == end - start + 1:
return path
if have > end - start + 1:
path.unlink()
have = 0
part = path.with_suffix(".part")
part.unlink(missing_ok=True)
subprocess.run(["curl", "-fsSL", "--connect-timeout", "12", "--max-time", "55",
"-r", f"{start + have}-{end}", "-o", str(part),
url + f"&range={start + have}"], check=False)
if part.is_file():
with path.open("ab") as stream:
stream.write(part.read_bytes())
if path.exists() and path.stat().st_size == end - start + 1:
return path
raise RuntimeError(f"Incomplete range {start}-{end}")
with ThreadPoolExecutor(16) as pool:
parts = list(pool.map(fetch, range(0, size, block)))
digest = hashlib.sha256()
output = root / "weights/v2m.mirror"
with output.open("wb") as stream:
for part in parts:
data = part.read_bytes()
digest.update(data)
stream.write(data)
if digest.hexdigest() != expected:
raise RuntimeError("Checkpoint checksum differs from the publisher's")
output.replace(root / "weights/yolo_world_v2m.pth")
print(digest.hexdigest(), flush=True)
|