| |
| |
| import argparse |
| import importlib.util |
| import json |
| import platform |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| def parser(): |
| p = argparse.ArgumentParser(description="Local image-to-Gaussian student research prototype") |
| sub = p.add_subparsers(dest="command", required=True) |
| sub.add_parser("doctor", help="Check dependencies, Metal availability and disk space; no downloads") |
| demo = sub.add_parser("demo-data", help="Create synthetic code-test fixtures, NOT teacher data") |
| demo.add_argument("--out", required=True) |
| demo.add_argument("--count", type=int, default=12) |
| demo.add_argument("--grid", type=int, choices=[8,16,32], default=8) |
| cache = sub.add_parser("cache", help="Cache Idol PLY with optional aligned geometry teacher meshes") |
| cache.add_argument("--image", required=True, help="Exact prepared/framed input used by teacher") |
| cache.add_argument("--idol-ply", required=True) |
| cache.add_argument("--out", required=True) |
| cache.add_argument("--id", required=True) |
| cache.add_argument("--object-id", help="Shared ID for all images of the same object, to prevent validation leakage") |
| cache.add_argument("--grid", type=int, choices=[8,16,32], default=16) |
| cache.add_argument("--coarse-mesh", help="InstantMesh output; this command imports, does not run InstantMesh") |
| cache.add_argument("--coarse-transform", help="JSON 4x4 mesh-to-Idol-PLY matrix") |
| cache.add_argument("--fine-mesh", help="Hi3D output; this command imports, does not run Hi3D") |
| cache.add_argument("--fine-transform", help="JSON 4x4 mesh-to-Idol-PLY matrix") |
| train = sub.add_parser("train", help="Train the student on cached targets") |
| train.add_argument("--manifest", required=True) |
| train.add_argument("--out", required=True) |
| train.add_argument("--device", choices=["auto","cuda","mps","cpu"], default="auto") |
| train.add_argument("--init-from", help="Post-train checkpoint weights; optimizer starts fresh") |
| train.add_argument("--steps", type=int, default=1000) |
| train.add_argument("--batch-size", type=int, default=1) |
| train.add_argument("--width", type=int, default=128) |
| train.add_argument("--lr", type=float, default=3e-4) |
| train.add_argument("--seed", type=int, default=42) |
| train.add_argument("--eval-every", type=int, default=50) |
| train.add_argument("--validation-fraction", type=float, default=.2) |
| train.add_argument("--allow-synthetic", action="store_true") |
| generate = sub.add_parser("generate", help="Generate standard 3DGS PLY with a trained student") |
| generate.add_argument("--checkpoint", required=True) |
| generate.add_argument("--image", required=True, help="Centered cutout or teacher-prepared image, black background") |
| generate.add_argument("--out", required=True) |
| generate.add_argument("--device", choices=["auto","cuda","mps","cpu"], default="auto") |
| generate.add_argument("--threshold", type=float, default=.5) |
| generate.add_argument("--max-splats", type=int, default=32768) |
| render = sub.add_parser("preview", help="CPU diagnostic turntable contact sheet") |
| render.add_argument("--ply", required=True) |
| render.add_argument("--out", required=True) |
| teacher = sub.add_parser("teacher-idol", help="EXPERIMENTAL staged Mac runner; weights not bundled") |
| teacher.add_argument("--image", required=True) |
| teacher.add_argument("--out", required=True) |
| teacher.add_argument("--cache", default=".cache/huggingface") |
| teacher.add_argument("--weights", default="speridlabs/idol") |
| teacher.add_argument("--device", choices=["mps","cpu"], default="mps") |
| teacher.add_argument("--mode", choices=["real","synthetic"], default="synthetic") |
| teacher.add_argument("--steps", type=int, default=20) |
| teacher.add_argument("--gaussians", type=int, default=32768) |
| teacher.add_argument("--query-chunk", type=int, default=128) |
| teacher.add_argument("--seed", type=int, default=42) |
| teacher.add_argument("--allow-download", action="store_true", help="Explicitly allow downloading model weights and gated encoders") |
| return p |
|
|
|
|
| def main(): |
| p = parser() |
| args = p.parse_args() |
| try: |
| if args.command == "doctor": |
| report = {"python": platform.python_version(), "platform": platform.platform(), "torch": str(torch.__version__), |
| "mps_available": torch.backends.mps.is_available(), "mps_built": torch.backends.mps.is_built(), |
| "free_disk_gib": round(shutil.disk_usage(Path.cwd()).free / 1024**3, 2), |
| "idol_installed": importlib.util.find_spec("idol") is not None, |
| "teacher_status": {"Idol": "experimental adapter, real weights unverified", |
| "InstantMesh": "cached mesh import only; no Mac execution adapter", |
| "Hi3D": "cached mesh import only; no Mac execution adapter"}, |
| "weights_bundled": False, "quantization": "FP32 student; no binary kernels or 1-bit claims"} |
| print(json.dumps(report, indent=2)) |
| elif args.command == "demo-data": |
| from .demo import create_demo |
| print(create_demo(args.out, args.count, args.grid)) |
| elif args.command == "cache": |
| from .data import cache_sample |
| from .splats import read_ply |
| print(json.dumps(cache_sample(args.out, args.id, args.image, read_ply(args.idol_ply), args.grid, |
| coarse=args.coarse_mesh, coarse_transform=args.coarse_transform, |
| fine=args.fine_mesh, fine_transform=args.fine_transform, |
| extra={"object_id": args.object_id or args.id, "idol_ply": str(Path(args.idol_ply).resolve())}), indent=2)) |
| elif args.command == "train": |
| from .engine import train |
| train(args) |
| elif args.command == "generate": |
| from .engine import generate |
| print(json.dumps(generate(args.checkpoint, args.image, args.out, args.device, args.threshold, args.max_splats), indent=2)) |
| elif args.command == "preview": |
| from PIL import Image, ImageDraw |
| import numpy as np |
| from .splats import preview, read_ply |
| splats = read_ply(args.ply) |
| sheet = Image.new("RGB", (1024, 280), "#111111") |
| draw = ImageDraw.Draw(sheet) |
| for i in range(4): |
| sheet.paste(preview(splats, yaw=i*np.pi/2), (i*256, 24)) |
| draw.text((i*256+8, 6), f"{i*90} degrees - diagnostic view", fill="white") |
| out = Path(args.out) |
| if out.exists(): |
| raise ValueError("Preview already exists") |
| out.parent.mkdir(parents=True, exist_ok=True) |
| sheet.save(out) |
| elif args.command == "teacher-idol": |
| from .teacher_idol import run_idol |
| run_idol(args) |
| except (ValueError, RuntimeError, FileNotFoundError, ImportError) as exc: |
| print(f"Error: {exc}", file=sys.stderr) |
| return 2 |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
