Upload encoder/code/k3_sheet.py with huggingface_hub

encoder/code/k3_sheet.py

@@ -0,0 +1,197 @@
+"""K3 sheet - four renders per photo, so quality is looked at instead of argued about.
+
+The acceptance test we never ran: put the teacher's 20-step result next to what our
+chain produces for the same photo, and look. Numbers like "lat cos 0.84" are a proxy
+for this image, and they were made on cosplay studio photos, not on the photos people
+actually drop in.
+
+Four PLYs per input, all through the teacher's own decoder so the only differences are
+the ones we care about:
+
+    t20      teacher conditions, 20 steps, its own sampler (guidance 3, shift 3)  <- truth
+    t2       teacher conditions,  2 steps, guidance 3                          <- cost of fewer steps
+    s2       student conditions,  2 steps, guidance 3                          <- cost of our encoder
+    s2g      student conditions,  2 steps, guidance 1, shift 1                 <- what the browser runs
+
+Also written per image: the canvas the student encoder saw (1024^2, black), the alpha
+the student matte produced, and the teacher's own prepared image. When a sheet cell is
+fog, one of those three says which half broke.
+
+Run it from the k3 driver (stage "sheet"), or directly:
+    FLIP_WORK=/kaggle/working python3 k3_sheet.py /kaggle/working/eval
+"""
+
+import importlib.util
+import json
+import os
+import pathlib
+import sys
+import tarfile
+import time
+
+WORK = pathlib.Path(os.environ.get("FLIP_WORK", "/kaggle/working"))
+K3 = WORK / "k3_encoder_distill.py"
+IMGDIR = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else str(WORK / "eval"))
+OUT = WORK / "k3" / "sheet"
+OUT.mkdir(parents=True, exist_ok=True)
+N_GAUSS = 65536
+
+os.environ.setdefault("FLIP_WORK", str(WORK))
+
+
+def p(*a):
+    print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)
+
+
+def load_k3():
+    spec = importlib.util.spec_from_file_location("k3", str(K3))
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def pad2(f):
+    import torch  # noqa: PLC0415
+    z = torch.zeros(f.shape[0], 5, f.shape[2], dtype=f.dtype, device=f.device)
+    return torch.cat([z, f], dim=1)
+
+
+def stats(g):
+    return {
+        "n": int(g.get_xyz.shape[0]),
+        "opacity": round(float(g.get_opacity.detach().float().mean()), 4),
+        "scale": round(float(g.get_scaling.detach().float().mean()), 5),
+        "x": [round(float(g.get_xyz[:, 0].min()), 3),
+              round(float(g.get_xyz[:, 0].max()), 3)],
+    }
+
+
+def main():
+    import numpy as np  # noqa: PLC0415
+    import torch  # noqa: PLC0415
+    import torch.nn.functional as F  # noqa: PLC0415
+    from PIL import Image  # noqa: PLC0415
+    from torchvision import transforms  # noqa: PLC0415
+
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    p("device", dev, "torch", torch.__version__)
+    k3 = load_k3()
+    sys.path.insert(0, str(k3.REPO))
+    import triposplat as T  # noqa: PLC0415
+
+    imgs = sorted(q for q in IMGDIR.rglob("*")
+                  if q.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp"))
+    limit = int(os.environ.get("FLIP_SHEET_N", "0"))
+    if limit:
+        imgs = imgs[:limit]
+    if not imgs:
+        raise SystemExit(f"no images under {IMGDIR}")
+    p("images", len(imgs), [q.name for q in imgs])
+
+    t0 = time.time()
+    pipe = k3.build_teacher_pipeline(T, dev)
+    p("teacher pipeline", round(time.time() - t0, 1), "s")
+    dino_s = k3.load_weights(k3.build_dino(), "dino", dev)
+    vae_s = k3.load_weights(k3.build_vae(), "vae", dev)
+    matte_s = k3.load_weights(k3.build_matte(), "matte", dev)
+    p("student encoders loaded")
+
+    report = []
+    for i, path in enumerate(imgs):
+        stem = path.stem
+        row = {"image": path.name, "stem": stem, "runs": {}}
+        try:
+            t_img = time.time()
+            image = Image.open(path).convert("RGB")
+            row["size"] = list(image.size)
+
+            # --- teacher path: its own preprocess, its own encoder
+            prepared = pipe.preprocess_image(image)
+            prepared.save(OUT / f"{stem}_prepared.png")
+            row["prepared"] = list(prepared.size)
+            gen = torch.Generator(device=dev).manual_seed(42)
+            cond_t = pipe.encode_image(prepared, generator=gen)
+
+            # --- student path: mirror of site/src/pipeline.js loadCondition
+            rgb = k3.resize_min_side(image)
+            w, h = rgb.size
+            xin = transforms.ToTensor()(rgb).unsqueeze(0).to(dev)
+            with torch.no_grad():
+                a = matte_s(xin)[0, 0].float().cpu().numpy()
+            alpha = (np.clip(a, 0, 1) * 255).astype(np.uint8)
+            Image.fromarray(alpha, "L").save(OUT / f"{stem}_mask.png")
+            er = np.array(k3.erode_alpha(alpha, 1), dtype=np.uint8)
+            box = k3.square_crop_box(er)
+            rgba = rgb.convert("RGBA")
+            rgba.putalpha(Image.fromarray(alpha, "L"))
+            cropped = rgba.crop(box).resize((k3.CANVAS, k3.CANVAS), Image.LANCZOS)
+            canvas = Image.new("RGB", (k3.CANVAS, k3.CANVAS), (0, 0, 0))
+            canvas.paste(cropped, mask=cropped.split()[3])
+            canvas.save(OUT / f"{stem}_canvas.png")
+            fx = alpha.astype(np.float32) / 255.0
+            row["mask"] = {
+                "resized": [w, h],
+                "box": [int(v) for v in box],
+                "fill": round(float(fx.mean()), 4),
+                "coverage": round(float((fx > 0.5).mean()), 4),
+            }
+
+            x = transforms.ToTensor()(canvas).unsqueeze(0).to(dev)
+            dino_dev = next(pipe.dinov3.parameters()).device
+            dino_dt = next(pipe.dinov3.parameters()).dtype
+            vae_dev = next(pipe.vae_encoder.parameters()).device
+            vae_dt = next(pipe.vae_encoder.parameters()).dtype
+            with torch.no_grad():
+                sf1 = dino_s(x)
+                sf2 = vae_s(x * 2 - 1)
+                # teacher features on the SAME canvas: isolates encoder error from crop
+                tf1 = F.layer_norm(pipe.dinov3(
+                    pixel_values=T._DINOV3_NORMALIZE(x).to(dino_dev, dino_dt)
+                ).float(), (1280,))
+                tf2 = pipe.vae_encoder.encode(
+                    (x * 2 - 1).to(vae_dev, vae_dt), deterministic=True)
+            row["cos1"] = round(float(F.cosine_similarity(sf1, tf1, dim=-1).mean()), 4)
+            row["cos2"] = round(float(F.cosine_similarity(sf2, tf2, dim=-1).mean()), 4)
+            p(stem, "cos", row["cos1"], row["cos2"], "mask",
+              row["mask"]["coverage"], "box", row["mask"]["box"])
+
+            cond_s = {"feature1": sf1, "feature2": pad2(sf2)}
+            variants = [
+                ("t20", cond_t, 20, 3.0, 3.0),
+                ("t2", cond_t, 2, 3.0, 3.0),
+                ("s2", cond_s, 2, 3.0, 3.0),
+                ("s2g", cond_s, 2, 1.0, 1.0),
+            ]
+            for tag, cond, steps, gs, shift in variants:
+                t1 = time.time()
+                gen = torch.Generator(device=dev).manual_seed(42)
+                lat = pipe.sample_latent(cond, steps=steps, guidance_scale=gs,
+                                         shift=shift, generator=gen)
+                g = pipe.decode_latent(lat["latent"], num_gaussians=N_GAUSS)
+                ply = OUT / f"{stem}_{tag}.ply"
+                g.save_ply(str(ply))
+                row["runs"][tag] = {
+                    **stats(g), "ply": ply.name,
+                    "mb": round(ply.stat().st_size / 1e6, 2),
+                    "secs": round(time.time() - t1, 1),
+                }
+                p(f"  {tag:4s} {row['runs'][tag]['secs']:6.1f}s "
+                  f"op {row['runs'][tag]['opacity']} sc {row['runs'][tag]['scale']}")
+            row["secs"] = round(time.time() - t_img, 1)
+        except Exception as exc:  # noqa: BLE001
+            import traceback  # noqa: PLC0415
+            row["error"] = traceback.format_exc()[-2000:]
+            p(stem, "FAILED", type(exc).__name__, str(exc)[:200])
+        report.append(row)
+        (OUT / "report.json").write_text(json.dumps(report, indent=2))
+
+    tgz = OUT / "sheet.tgz"
+    with tarfile.open(tgz, "w:gz") as tf:
+        for f in sorted(OUT.iterdir()):
+            if f.name != tgz.name:
+                tf.add(f, arcname=f.name)
+    p("wrote", tgz, round(tgz.stat().st_size / 1e6, 1), "MB")
+
+
+if __name__ == "__main__":
+    main()