| """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. |
| |
| 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), FLIP_HF_BASE (force one host). |
| """ |
|
|
| import importlib.util |
| import os |
| import pathlib |
| import shutil |
| import subprocess |
| import sys |
| import tarfile |
| import time |
| import urllib.request |
|
|
| 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") |
| BASES = [b for b in (os.environ.get("FLIP_HF_BASE"), |
| "https://huggingface.co", "https://hf-mirror.com") if b] |
|
|
| os.environ.setdefault("PYTHONUNBUFFERED", "1") |
| if not os.environ.get("FLIP_HF_BASE"): |
| |
| |
| os.environ.pop("HF_ENDPOINT", None) |
|
|
|
|
| def p(*a): |
| print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True) |
|
|
|
|
| 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: |
| 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: |
| 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) |
| fetch("encoder/code/k3_encoder_distill.py", K3) |
| fetch("encoder/code/k3_sheet.py", WORK / "k3_sheet.py") |
| if not (WORK / "inputs").is_dir(): |
| 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")): |
| 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) |
|
|
|
|
| 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, "py", sys.version.split()[0]) |
| probe() |
| 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("matte: resume from weights/matte.pt, target epoch", EPOCHS) |
| k3.train_matte("cuda") |
| import torch |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| if "export" in STAGES: |
| k3 = k3 or load_k3() |
| p("export") |
| p("meta", k3.export()) |
| if "e2e" in STAGES: |
| k3 = k3 or load_k3() |
| p("e2e") |
| k3.e2e() |
| p((OUT / "e2e.json").read_text()) |
| if "sheet" in STAGES: |
| |
| |
| env = dict(os.environ, FLIP_WORK=str(WORK)) |
| src = os.environ.get("FLIP_SHEET_DIR", str(WORK / "eval")) |
| cmd = [sys.executable, str(WORK / "k3_sheet.py"), src] |
| p("$", " ".join(cmd)) |
| subprocess.run(cmd, env=env, check=True) |
| 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) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
