colibri/c/download_fp8.py

"""Download GLM-5.2-FP8 from ModelScope (fast, no HF throttling) with
HuggingFace fallback. Parallel shard download, clean progress display.
Usage: python download_fp8.py
       python download_fp8.py --parallel 4
       python download_fp8.py --source hf  (force HuggingFace)
"""
import os, sys, time, threading, argparse, subprocess

REPO_MS = "ZhipuAI/GLM-5.2-FP8"        # ModelScope
REPO_HF = "zai-org/GLM-5.2-FP8"        # HuggingFace
DEST = r"I:\glm52_fp8"

# ── ANSI colors ──
class C:
    dim="\033[2m"; grn="\033[32m"; yel="\033[33m"; cyn="\033[36m"; b="\033[1m"; r="\033[0m"

def fmt_bytes(n):
    if n>=1e9: return f"{n/1e9:.2f} GB"
    if n>=1e6: return f"{n/1e6:.1f} MB"
    return f"{n/1e3:.0f} KB"

def fmt_time(s):
    if s<60: return f"{s:.0f}s"
    if s<3600: return f"{s/60:.0f}m"
    return f"{s/3600:.1f}h"

def bar(cur, total, width=24):
    if total<=0: return "["+" "*width+"]"
    pct=min(cur/total,1.0); filled=int(width*pct)
    return "["+"█"*filled+"░"*(width-filled)+f"] {pct*100:4.0f}%"

def get_shard_list_hf():
    from huggingface_hub import HfApi
    info=HfApi().repo_info(REPO_HF, files_metadata=True)
    shards=sorted(s.rfilename for s in info.siblings if s.rfilename.endswith(".safetensors"))
    sizes={s.rfilename:s.size for s in info.siblings if s.rfilename.endswith(".safetensors")}
    return shards, sizes

def get_shard_list_ms():
    """Get shard list from ModelScope API."""
    import requests
    # ModelScope API: list files
    r = requests.get(f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo/files?Revision=master&Root=", timeout=30)
    data = r.json()["Data"]["Files"]
    shards = sorted(f["Path"] for f in data if f["Path"].endswith(".safetensors"))
    sizes = {f["Path"]: f.get("Size", 0) for f in data if f["Path"].endswith(".safetensors")}
    return shards, sizes

def download_file_ms(fn):
    """Download a single file from ModelScope using their CDN."""
    from modelscope.hub.file_download import model_file_download
    model_file_download(
        model_id=REPO_MS,
        file_path=fn,
        local_dir=DEST,
        revision="master",
    )

def download_file_hf(fn):
    """Download a single file from HuggingFace with hf_transfer."""
    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
    from huggingface_hub import hf_hub_download
    hf_hub_download(REPO_HF, fn, local_dir=DEST)

def download_file_curl(fn, base_url, expected_size):
    """Fallback: download with curl to a .part file with resume."""
    outpath = os.path.join(DEST, fn)
    partpath = outpath + ".part"
    if os.path.exists(outpath) and os.path.getsize(outpath) == expected_size:
        return True
    url = f"{base_url}/{fn}"
    cmd = ["curl", "-L", "-C", "-", "--retry", "999", "--retry-delay", "5",
           "--connect-timeout", "15", "--speed-time", "30", "--speed-limit", "1000",
           "-o", partpath, "-H", "User-Agent: colibri-download/1.0", url]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if os.path.exists(partpath) and os.path.getsize(partpath) >= expected_size:
        os.replace(partpath, outpath)
        return True
    return False

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--parallel", type=int, default=3)
    ap.add_argument("--source", choices=["auto", "ms", "hf"], default="auto",
                    help="auto (try ModelScope first), ms (ModelScope only), hf (HuggingFace only)")
    args = ap.parse_args()

    os.makedirs(DEST, exist_ok=True)

    # Determine source and get shard list
    use_ms = False
    shards, sizes = [], {}

    if args.source in ("auto", "ms"):
        try:
            print(f"{C.dim}Trying ModelScope...{C.r}", end=" ", flush=True)
            shards, sizes = get_shard_list_ms()
            use_ms = True
            print(f"{C.grn}✓{C.r} {len(shards)} shards found")
        except Exception as e:
            print(f"{C.yel}failed ({e}){C.r}")
            if args.source == "ms":
                print("ModelScope failed and --source ms was set. Exiting."); return

    if not shards:
        print(f"{C.dim}Using HuggingFace...{C.r}", end=" ", flush=True)
        shards, sizes = get_shard_list_hf()
        print(f"{C.grn}✓{C.r} {len(shards)} shards found")

    total = len(shards)
    total_bytes = sum(sizes.values())
    source_name = "ModelScope" if use_ms else "HuggingFace"

    # Download metadata files
    meta_files = ["config.json", "tokenizer.json", "tokenizer_config.json",
                  "generation_config.json", "model.safetensors.index.json"]
    for fn in meta_files:
        out = os.path.join(DEST, fn)
        if not os.path.exists(out):
            try:
                if use_ms: download_file_ms(fn)
                else: download_file_hf(fn)
            except Exception: pass

    # Build work queue
    todo = []
    done_set = set()
    for fn in shards:
        outpath = os.path.join(DEST, fn)
        if os.path.exists(outpath) and os.path.getsize(outpath) == sizes.get(fn, 0):
            done_set.add(fn)
        else:
            todo.append(fn)

    existing = len(done_set)
    print(f"\n{C.b}GLM-5.2-FP8 Download ({source_name}){C.r}")
    print(f"  {C.dim}{total} shards · {total_bytes/1e9:.0f} GB · {existing}/{total} complete{C.r}")
    print(f"  {C.dim}{len(todo)} to download · {args.parallel} parallel{C.r}")
    if todo:
        remaining = sum(sizes[fn] for fn in todo)
        print(f"  {C.dim}Remaining: {remaining/1e9:.0f} GB{C.r}")
    print()

    if not todo:
        print(f"{C.grn}✓ All shards already downloaded!{C.r}\n"); return

    lock = threading.Lock()
    completed = list(done_set)
    t0 = time.time()
    qidx = [0]

    def worker(wid):
        while True:
            with lock:
                if qidx[0] >= len(todo): return
                idx = qidx[0]; qidx[0] += 1
                fn = todo[idx]

            expected = sizes.get(fn, 0)
            shard_num = existing + idx + 1
            print(f"  {C.cyn}[{shard_num}/{total}]{C.r} {C.dim}{fn}{C.r}")

            success = False
            for attempt in range(3):
                try:
                    if use_ms:
                        download_file_ms(fn)
                    else:
                        download_file_hf(fn)
                    outpath = os.path.join(DEST, fn)
                    # ModelScope downloads to a cache dir, need to check
                    # if the file exists at our expected path
                    if not os.path.exists(outpath):
                        # Try to find it in ModelScope's cache structure
                        # and copy/symlink it
                        pass
                    if os.path.exists(outpath):
                        actual = os.path.getsize(outpath)
                        if expected == 0 or actual == expected:
                            success = True; break
                        # Size mismatch — might be in cache
                    # If we got here, file downloaded but not at expected path
                    # ModelScope puts it in local_dir; check again
                    if os.path.exists(outpath):
                        success = True; break
                    # Retry with curl fallback
                    if use_ms:
                        base = f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo?Revision=master&FilePath="
                    else:
                        base = f"https://huggingface.co/{REPO_HF}/resolve/main"
                    if download_file_curl(fn, base, expected):
                        success = True; break
                except Exception as e:
                    if attempt < 2:
                        print(f"    {C.yel}retry {attempt+1}: {e}{C.r}")
                        time.sleep(3)
                    else:
                        print(f"    {C.yel}✗ failed: {e}{C.r}")

            with lock:
                if success:
                    completed.append(fn)
                    elapsed = time.time() - t0
                    have = sum(sizes.get(f,0) for f in completed)
                    pct = 100.0 * have / total_bytes
                    speed = (have - sum(sizes[f] for f in done_set)) / max(elapsed, 1)
                    eta = (total_bytes - have) / speed if speed > 0 else 0
                    print(f"  {C.grn}✓{C.r} {fn} {C.dim}— {len(completed)}/{total} · "
                          f"{pct:.1f}% · {fmt_time(elapsed)} · ETA {fmt_time(eta)}{C.r}")
                else:
                    print(f"  {C.yel}✗ GIVE UP: {fn}{C.r}")

    threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(args.parallel)]
    for t in threads: t.start()
    for t in threads: t.join()

    print()
    final = sum(1 for fn in shards
                if os.path.exists(os.path.join(DEST, fn))
                and os.path.getsize(os.path.join(DEST, fn)) == sizes.get(fn, 0))
    if final == total:
        print(f"{C.grn}{'='*50}")
        print(f"  ✓ All {total} shards downloaded!{C.r}\n")
    else:
        print(f"{C.yel}  {final}/{total} complete, {total-final} remaining{C.r}")
        print(f"  Re-run to resume.\n")

if __name__ == "__main__":
    try: main()
    except KeyboardInterrupt:
        print(f"\n\n{C.yel}Interrupted. Re-run to resume — no data lost.{C.r}\n")