Upload encoder/code/k3_kaggle_driver.py with huggingface_hub

encoder/code/k3_kaggle_driver.py

@@ -0,0 +1,117 @@
+"""K3 on Kaggle: the driver the notebook cell downloads and runs.
+
+The notebook itself is six lines (fetch this file from HF, run it). This file
+holds the sequence so a change is one upload, not one edited cell:
+
+    1. restore  inputs.tgz + the three checkpoints from idefinity/flip-student-2step
+    2. setup,prep   clone upstream, 4.46 GB teachers, cache 200 images
+    3. matte 7 -> 20  (reads k3/weights/matte.pt, which is matte_full.pt)
+    4. export  fp32 ONNX + parity, writes encoder_meta.json
+    5. e2e     teacher DiT 2-step, student vs teacher conditions
+
+Everything lands in /kaggle/working/k3, which Kaggle saves as the version output.
+
+Env: FLIP_NB_STAGES (default "restore,setup,prep,train,export,e2e"),
+     FLIP_EPOCHS (matte target, default 20), HF_ENDPOINT (default the mirror).
+"""
+
+import importlib.util
+import os
+import pathlib
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+
+WORK = pathlib.Path(os.environ.get("FLIP_WORK", "/kaggle/working"))
+OUT = WORK / "k3"
+REPO_ID = "idefinity/flip-student-2step"
+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")
+
+os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
+os.environ.setdefault("PYTHONUNBUFFERED", "1")
+
+
+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 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")
+    if not (WORK / "inputs").is_dir():
+        with tarfile.open(fetch("encoder/checkpoints/inputs.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)
+    imgs = sorted((WORK / "inputs").rglob("*.jp*g"))
+    p("inputs", len(imgs), "->", imgs[0].name if imgs else None)
+
+
+def run_k3(stages):
+    env = dict(os.environ, FLIP_WORK=str(WORK), FLIP_IMAGE_DIR=str(WORK / "inputs"),
+               FLIP_N_IMG="200", FLIP_EPOCHS=EPOCHS)
+    cmd = [sys.executable, str(K3), "--stages", stages]
+    p("$", " ".join(cmd))
+    subprocess.run(cmd, env=env, check=True)
+
+
+def load_k3():
+    os.environ.update(
+        FLIP_WORK=str(WORK), FLIP_DINO_ARCH="conv", FLIP_DINO_DIM="512",
+        FLIP_DINO_DEPTH="8", FLIP_FP16="0",
+        PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True")
+    spec = importlib.util.spec_from_file_location("k3", str(K3))
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    mod.EPOCHS = int(EPOCHS)
+    return mod
+
+
+def main():
+    p("stages", STAGES, "epochs", EPOCHS)
+    if "restore" in STAGES:
+        restore()
+    if "setup" in STAGES or "prep" in STAGES:
+        run_k3(",".join(s for s in STAGES if s in ("setup", "prep")))
+    k3 = None
+    if "train" in STAGES:
+        k3 = load_k3()
+        p("train matte on cuda; resumes from epoch in weights/matte.pt")
+        k3.train_matte("cuda")
+        import torch  # noqa: PLC0415
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    if "export" in STAGES:
+        k3 = k3 or load_k3()
+        p("export")
+        meta = k3.export()
+        p("meta", meta)
+    if "e2e" in STAGES:
+        k3 = k3 or load_k3()
+        p("e2e")
+        k3.e2e()
+        p((OUT / "e2e.json").read_text())
+    p("outputs")
+    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)
+
+
+if __name__ == "__main__":
+    main()