Upload encoder/code/k3_kaggle_driver.py with huggingface_hub

encoder/code/k3_kaggle_driver.py

@@ -11,8 +11,13 @@ holds the sequence so a change is one upload, not one edited cell:
 
 Everything lands in /kaggle/working/k3, which Kaggle saves as the version output.
 
+Why not huggingface_hub for the restore: resolving a file through hf-mirror
+loses the x-repo-commit header, and the hub client reads that as "not on
+huggingface.co". Plain GET on /resolve/main/<path> works on both hosts, so the
+restore goes that way and tries huggingface.co first, the mirror second.
+
 Env: FLIP_NB_STAGES (default "restore,setup,prep,train,export,e2e"),
-     FLIP_EPOCHS (matte target, default 20), HF_ENDPOINT (default the mirror).
+     FLIP_EPOCHS (matte target, default 20), FLIP_HF_BASE (force one host).
 """
 
 import importlib.util
@@ -23,6 +28,7 @@ import subprocess
 import sys
 import tarfile
 import time
+import urllib.request
 
 WORK = pathlib.Path(os.environ.get("FLIP_WORK", "/kaggle/working"))
 OUT = WORK / "k3"
@@ -31,8 +37,9 @@ K3 = WORK / "k3_encoder_distill.py"
 STAGES = [s.strip() for s in os.environ.get(
     "FLIP_NB_STAGES", "restore,setup,prep,train,export,e2e").split(",") if s.strip()]
 EPOCHS = os.environ.get("FLIP_EPOCHS", "20")
+BASES = [b for b in (os.environ.get("FLIP_HF_BASE"),
+                     "https://huggingface.co", "https://hf-mirror.com") if b]
 
-os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
 os.environ.setdefault("PYTHONUNBUFFERED", "1")
 
 
@@ -40,25 +47,59 @@ def p(*a):
     print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)
 
 
-def fetch(path):
-    from huggingface_hub import hf_hub_download  # noqa: PLC0415
-    p("fetch", path)
-    return hf_hub_download(REPO_ID, path, local_dir=str(WORK / "dl"))
+def probe():
+    for base in BASES:
+        t0 = time.time()
+        try:
+            req = urllib.request.Request(f"{base}/api/models/{REPO_ID}",
+                                         headers={"User-Agent": "flip-k3"})
+            with urllib.request.urlopen(req, timeout=25) as fh:
+                p("probe", base, fh.status, f"{time.time() - t0:.1f}s")
+        except Exception as exc:  # noqa: BLE001
+            p("probe", base, "FAIL", type(exc).__name__, str(exc)[:120])
+
+
+def get(url, timeout=600):
+    req = urllib.request.Request(url, headers={"User-Agent": "flip-k3"})
+    with urllib.request.urlopen(req, timeout=timeout) as fh:
+        return fh.read()
+
+
+def fetch(path, dst, timeout=600):
+    dst = pathlib.Path(dst)
+    if dst.exists() and dst.stat().st_size > 1024:
+        p("have", dst.name, dst.stat().st_size)
+        return dst
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    err = None
+    for base in BASES:
+        try:
+            t0 = time.time()
+            blob = get(f"{base}/{REPO_ID}/resolve/main/{path}", timeout)
+            tmp = dst.with_suffix(dst.suffix + ".part")
+            tmp.write_bytes(blob)
+            os.replace(tmp, dst)
+            p("got", dst.name, dst.stat().st_size,
+              f"{time.time() - t0:.0f}s from", base)
+            return dst
+        except Exception as exc:  # noqa: BLE001
+            err = f"{type(exc).__name__}: {str(exc)[:100]}"
+            p("  miss", base, err)
+    raise RuntimeError(f"cannot fetch {path}: {err}")
 
 
 def restore():
     W = OUT / "weights"
     W.mkdir(parents=True, exist_ok=True)
-    shutil.copy(fetch("encoder/code/k3_encoder_distill.py"), K3)
-    p("k3 script", K3.stat().st_size, "B")
+    fetch("encoder/code/k3_encoder_distill.py", K3)
     if not (WORK / "inputs").is_dir():
-        with tarfile.open(fetch("encoder/checkpoints/inputs.tgz")) as tf:
+        tgz = fetch("encoder/checkpoints/inputs.tgz", WORK / "dl/inputs.tgz")
+        with tarfile.open(tgz) as tf:
             tf.extractall(WORK)
     for src, dst in (("matte_full.pt", "matte.pt"),
                      ("dino_model.pt", "dino.pt"),
                      ("vae_model.pt", "vae.pt")):
-        shutil.copy(fetch(f"encoder/checkpoints/{src}"), W / dst)
-        p("weights", dst, (W / dst).stat().st_size)
+        fetch(f"encoder/checkpoints/{src}", W / dst)
     imgs = sorted((WORK / "inputs").rglob("*.jp*g"))
     p("inputs", len(imgs), "->", imgs[0].name if imgs else None)
 
@@ -84,7 +125,8 @@ def load_k3():
 
 
 def main():
-    p("stages", STAGES, "epochs", EPOCHS)
+    p("stages", STAGES, "epochs", EPOCHS, "py", sys.version.split()[0])
+    probe()
     if "restore" in STAGES:
         restore()
     if "setup" in STAGES or "prep" in STAGES:
@@ -92,7 +134,7 @@ def main():
     k3 = None
     if "train" in STAGES:
         k3 = load_k3()
-        p("train matte on cuda; resumes from epoch in weights/matte.pt")
+        p("matte: resume from weights/matte.pt, target epoch", EPOCHS)
         k3.train_matte("cuda")
         import torch  # noqa: PLC0415
         if torch.cuda.is_available():
@@ -100,14 +142,13 @@ def main():
     if "export" in STAGES:
         k3 = k3 or load_k3()
         p("export")
-        meta = k3.export()
-        p("meta", meta)
+        p("meta", k3.export())
     if "e2e" in STAGES:
         k3 = k3 or load_k3()
         p("e2e")
         k3.e2e()
         p((OUT / "e2e.json").read_text())
-    p("outputs")
+    p("outputs under", OUT)
     for f in sorted(OUT.rglob("*")):
         if f.is_file() and f.suffix in (".onnx", ".json", ".pt", ".txt"):
             p(" ", f.relative_to(WORK), f.stat().st_size)