| |
| """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() |
|
|