File size: 5,616 Bytes
09cf3ce | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | #!/usr/bin/env python3
"""Mirror the ACE-Step production model packages from the upstream R2 bucket.
Downloads the three content-addressed packages (direct reference, DiT rev7,
VAE rev7), verifies every payload's SHA-256 and byte length against its
manifest, and checks the aggregate inventory against the counts documented in
the upstream demo README (113 payloads + 3 manifests, 5,749,459,255 bytes).
Files the reference manifest lists but the direct R2 deployment deliberately
excludes (planner/semantic material, VAE-only shards, superseded DiT layers)
404 and are recorded as `excluded` — the success criterion is byte-exact
match of what IS served, not of the full manifest.
"""
import concurrent.futures as cf
import hashlib
import json
import os
import sys
import urllib.request
ORIGIN = "https://ace-step-wgsl-models.narcotic.sh"
DEST = os.path.dirname(os.path.abspath(__file__))
PACKAGES = [
{
"id": "reference-direct",
"prefix": "v1/reference/18f36c6420976475af65ecd833ca56c6119706322ce54120389d4915d8e80db6",
"manifest": "direct-manifest-b44a3d157009d035a8f20aa752db4ceef2fac5bd140eff13be8f7488bc978089.json",
"manifest_sha256": "b44a3d157009d035a8f20aa752db4ceef2fac5bd140eff13be8f7488bc978089",
"expect_files": 58,
"expect_bytes": 2_558_130_431,
},
{
"id": "dit-revision7",
"prefix": "v1/dit-revision7/d3fc0020efcf60702db411da2fd4b93e9bb84f1437ed310aef01c892727e452f",
"manifest": "manifest.json",
"manifest_sha256": "d3fc0020efcf60702db411da2fd4b93e9bb84f1437ed310aef01c892727e452f",
"expect_files": 48,
"expect_bytes": 3_020_808_192,
},
{
"id": "vae-revision7",
"prefix": "v1/vae-revision7/36a54d79777d6826088095ba6ebc028fb4bea546368c0f0a29cd0eee8d656da7",
"manifest": "manifest.json",
"manifest_sha256": "36a54d79777d6826088095ba6ebc028fb4bea546368c0f0a29cd0eee8d656da7",
"expect_files": 7,
"expect_bytes": 168_791_552,
},
]
def fetch(url: str) -> bytes:
req = urllib.request.Request(url, headers={"User-Agent": "fluidaudio-mirror/1"})
with urllib.request.urlopen(req, timeout=120) as r:
return r.read()
def mirror_file(prefix: str, name: str, sha256: str, byte_length: int):
"""Download one payload; returns (name, status, bytes)."""
out = os.path.join(DEST, prefix, name)
if os.path.exists(out):
data = open(out, "rb").read()
if len(data) == byte_length and hashlib.sha256(data).hexdigest() == sha256:
return name, "cached", byte_length
try:
data = fetch(f"{ORIGIN}/{prefix}/{name}")
except urllib.error.HTTPError as e:
if e.code == 404:
return name, "excluded", 0
raise
digest = hashlib.sha256(data).hexdigest()
if digest != sha256:
return name, f"HASH MISMATCH {digest}", len(data)
if len(data) != byte_length:
return name, f"SIZE MISMATCH {len(data)}", len(data)
os.makedirs(os.path.dirname(out), exist_ok=True)
with open(out + ".tmp", "wb") as f:
f.write(data)
os.replace(out + ".tmp", out)
return name, "ok", byte_length
def main():
grand_files = 0
grand_bytes = 0
failures = []
for pkg in PACKAGES:
murl = f"{ORIGIN}/{pkg['prefix']}/{pkg['manifest']}"
mbytes = fetch(murl)
mdigest = hashlib.sha256(mbytes).hexdigest()
if mdigest != pkg["manifest_sha256"]:
print(f"FATAL: manifest hash mismatch for {pkg['id']}: {mdigest}")
sys.exit(1)
mpath = os.path.join(DEST, pkg["prefix"], pkg["manifest"])
os.makedirs(os.path.dirname(mpath), exist_ok=True)
open(mpath, "wb").write(mbytes)
manifest = json.loads(mbytes)
files = manifest["files"]
print(f"[{pkg['id']}] manifest verified, {len(files)} file records", flush=True)
got_files = 0
got_bytes = 0
excluded = 0
with cf.ThreadPoolExecutor(max_workers=6) as ex:
futures = [
ex.submit(mirror_file, pkg["prefix"], f["name"], f["sha256"], f["byteLength"])
for f in files
]
for i, fut in enumerate(cf.as_completed(futures)):
name, status, nbytes = fut.result()
if status in ("ok", "cached"):
got_files += 1
got_bytes += nbytes
elif status == "excluded":
excluded += 1
else:
failures.append((pkg["id"], name, status))
if (i + 1) % 20 == 0:
print(f"[{pkg['id']}] {i + 1}/{len(files)} processed, {got_bytes/1e9:.2f} GB", flush=True)
ok_count = got_files == pkg["expect_files"] and got_bytes == pkg["expect_bytes"]
print(
f"[{pkg['id']}] {'PASS' if ok_count else 'MISMATCH'}: "
f"{got_files} files / {got_bytes} bytes "
f"(expected {pkg['expect_files']} / {pkg['expect_bytes']}), {excluded} excluded-by-design",
flush=True,
)
if not ok_count:
failures.append((pkg["id"], "<inventory>", f"{got_files} files / {got_bytes} bytes"))
grand_files += got_files
grand_bytes += got_bytes
print(f"TOTAL: {grand_files} payloads, {grand_bytes} bytes (expected 113 / 5747730175)")
if failures:
print("FAILURES:")
for pkg_id, name, status in failures:
print(f" {pkg_id}: {name}: {status}")
sys.exit(1)
print("MIRROR COMPLETE AND VERIFIED")
if __name__ == "__main__":
main()
|