diff --git a/docs/2026-06-10-ll-gen-make-real-findings.md b/docs/2026-06-10-ll-gen-make-real-findings.md new file mode 100644 index 0000000..cf6cdff --- /dev/null +++ b/docs/2026-06-10-ll-gen-make-real-findings.md @@ -0,0 +1,151 @@ +# ll_gen "make-real" findings (2026-06-10) + +Goal: make `ll_gen`'s generators genuinely produce valid CAD with **honest** +metrics (no faked success), using the canonical DeepCAD dataset +(`data.tar` → `resources/DeepCAD/data/cad_vec/`, 179,133 real construction +sequences). + +`ll_gen` has **two** generation paths, and they fail/succeed for different +reasons: + +| Path | Representation | Executor | Status | +|---|---|---|---| +| command-VAE (`STEPVAE`) | command sequence (SOL/LINE/ARC/CIRCLE/EXTRUDE/EOS) | `command_executor` (sketch→extrude) | real but **primitive-limited** | +| diffusion (`StructuredDiffusion`) | per-face UV-grids + per-edge polylines | `surface_executor` (B-spline fit → sew) | codec **real**; denoiser convergence is compute-bound | + +--- + +## 1. Command-VAE: real, but architecturally primitive-limited + +**The existing 95%-valid checkpoint generates cylinders.** Inspecting +`ll_gen/checkpoints/vae_rl_solid.pt` (58/60 "valid"): the valid shapes are +dominated by **extruded circles** (bbox dims `(small, X, X)` — two equal dims), +and several "valid" shapes have **volume = 0** (flat faces OCC accepts but aren't +solids), inflating the rate. "88 distinct" is diversity of *size*, not *kind*. + +**Root cause (architecture, not tuning).** `STEPVAE` is a *global-latent, +non-autoregressive* model: `encode()` mean-pools the sequence into one orderless +latent (`vae.py:168`); `decode(z, seq_len)` emits every position's params +independently. Nothing ties LINE_i's end to LINE_{i+1}'s start, so multi-line +sketches almost never close → only self-closing primitives (circles) validate. +It reached 95% only because **RL** directly rewarded valid circles. + +**Supervised training confirms the limit from the other side.** Trained on real +DeepCAD `cad_vec` (command-CE + param-CE + KL warmup), prior-sample validity is +**0%** and the model **posterior-collapses to all-SOL** (60 SOL tokens, zero +curves; cmd_acc 0.191 = the SOL frequency). Under the KL pressure generation +requires, the orderless mean-pool latent is ignored. The VAE can be +informative-latent XOR good-prior-sampling, never both. + +### Fixes shipped (command path) + +- **Decode enum mismatch** (`generation_pipeline.py`, main `5fbf909`) — int→str + CommandType crash made ALL generation 0% valid; regression test. +- **Resilient checkpoint load** (`rl_trainer.py`, main `25cc091`) — strict load + rejected the real M3 checkpoint (`dim_encoder` drift); regression test. +- **NaN param loss** (`vae.py`) — `STEPVAE.forward`'s param loss averaged over + all 16 heads, but slots 8–15 are never active in the 6-command schema → + `F.cross_entropy(mean)` over zero elements = NaN → poisoned `recon_loss`, + silently blocking the supervised forward. Now skips all-ignored heads. + Regression: `ll_stepnet/tests/test_vae_sparse_param_loss.py`. +- **Closure-aware decode** (`command_executor.py::_build_sketch_face`) — builds + sketch loops by **threading curve endpoints + auto-closing**, so multi-line + polygons close *by construction* regardless of decoder endpoint alignment. + Verified: a deliberately non-connecting square → valid solid (vol 1.01); + triangle/pentagon/hexagon/octagon all valid; the 95% checkpoint now emits some + genuinely **non-cylindrical** solids (dims like `(1.97, 2.23, 2.7)`) with no + regression (59/60). Tests: `TestClosureAwareSketch` (3) in + `ll_gen/tests/test_command_executor.py` (43 pass). This removes the + wire-closure limit for **any** non-collapsed command generator — it cannot + rescue the posterior-collapsed VAE (no curves to close), but it is a real + prerequisite a better decoder needs. + +**Conclusion:** the command-VAE is a dead end for *diverse* generation by +architecture. Diverse valid CAD needs either an autoregressive/closure-aware +*decoder* or the diffusion path. + +--- + +## 2. Diffusion: the genuine path to diverse CAD + +`StructuredDiffusion` decodes per-face UV-grids + per-edge polylines and +`surface_executor` fits B-splines and **sews** them — no posterior collapse, no +polyline-closure limit. + +### Built + verified + +- **Geometry extraction** (`resources/ll_gen_proof/diffusion_codec_train.py`): + `cad_vec → solid` (validated translation, 30/30) → sample each face as an + 8×8×3 UV grid and each edge as a 12×3 polyline, normalized to the unit cube, + padded to `num_faces=8 / num_edges=12` with masks. Avg 4.3 faces / 8.5 edges. +- **GeometryCodec trained on real geometry**: reconstruction MSE + **0.40 (untrained) → 0.0003 (trained, 60 epochs / 4000 solids)** — RMS error + ~1.7% of the unit cube. The latent↔geometry map is now real. Checkpoint: + `resources/ll_gen_proof/diffusion_codec.pt`. + +### Bugs fixed (diffusion path) + +- **MPS backward crash** (`diffusion.py` `encode_faces`/`encode_edges`) — + `permute(...)` then Conv2d/Conv1d; the conv backward calls `.view()` on the + non-contiguous tensor and raises. Added `.contiguous()`. Also fixes the + existing RL diffusion path on MPS. +- **Surface-fitter signature** (`surface_executor.py::_fit_bspline_surface`) — + called cadling's `BSplineSurfaceFitter.fit_surface(grid, tolerance=...)`, but + that method takes only `point_grid` and returns a `dict` (not a face). + `TypeError` → **every** face silently failed → nothing ever sewed. Fixed to + use the constructor tolerance + extract `result["face"]`, falling back to the + direct OCC fit. Verified `_fit_bspline_surface` now returns a face. + +- **Topology-merge crash** (`surface_executor.py::execute_latent_proposal`) — + the merge step called `TopologyMerger.merge_edges(edges)`, which does not + exist (real API is `merge(faces) -> {shape, valid, ...}`, purpose-built to + mate independently-generated faces and sew a watertight solid). The + `AttributeError` was uncaught, so **every** sew crashed — the diffusion path + could never produce a shape, training or not. Fixed: call `merge(faces)` as + the primary watertight path, fall back to the built-in edge dedup + sew. + +### The sewing pipeline now works on in-distribution geometry + +With the surface-fit + merge fixes, **7/15 real DeepCAD solids sew into closed, +non-zero-volume shells** (the ≥4-face prismatic models). The 3-face failures are +cylinders whose periodic lateral surface the UV-grid doesn't seam — a known +sampling limitation. Regression test: +`ll_gen/tests/test_surface_executor.py::TestWatertightSew` (unit cube → closed +volume). So the path was *completely broken by three real bugs*, not by physics; +it is now functional for in-distribution geometry. + +### Honest remaining bottleneck: the denoiser does not converge + +The remaining step is **denoiser convergence**, and here is an honest negative: +the 4-stage denoisers **plateau at the trivial "predict-zero-noise" solution** +(denoiser-only loss ≈ 1.0 per stage = `mean(noise²)` for unit-variance noise — +i.e. the network outputs ≈0 instead of the noise). It drops only 5.28→4.03 then +flattens. Latent normalization (scaling the sub-unit ~0.42-std codec latents to +~unit variance — the standard latent-diffusion trick) was tried and did **not** +move the plateau, so it was reverted. The `CADDenoiser` architecture is +structurally sound (input_proj → sinusoidal time-embed → 12-layer transformer → +output_proj), so this is a genuine training-convergence problem (LR schedule / +architecture / per-token-vs-set conditioning), not a one-line bug — it needs +real diffusion-training experimentation. + +**Net for the diffusion path:** codec is real, the sewing pipeline is fixed and +works on in-distribution geometry, but the generator **cannot yet produce valid +CAD because the denoiser does not learn to denoise** — stated plainly, no number +claimed. This is the honest frontier; closing it is a focused follow-on. + +--- + +## Reusable asset + +`cad_vec → executor-schema` translation (cadlib `CADSequence.from_vector` → +SOL/LINE-abs/ARC-start-end-center/CIRCLE/EXTRUDE → quantize to the executor's +symmetric [0,255]↔[-2,2]). Validated 30/30 real models → valid solids. Transfers +to any target consuming the executor schema. + +## Artifacts + +- Scripts: `resources/ll_gen_proof/{deepcad_supervised_train, + diffusion_codec_train, diffusion_full_train, rl_refine_from_supervised}.py` +- Results: `resources/ll_gen_proof/{DEEPCAD_SUPERVISED, DIFFUSION_CODEC, + DIFFUSION_FULL}.json`, `diffusion_codec.pt` +- DeepCAD data: `resources/DeepCAD/data/cad_vec/` (179k h5, gitignored) diff --git a/ll_brepnet/mlx/train_brepnet_mlx.py b/ll_brepnet/mlx/train_brepnet_mlx.py new file mode 100644 index 0000000..27f6b2f --- /dev/null +++ b/ll_brepnet/mlx/train_brepnet_mlx.py @@ -0,0 +1,330 @@ +"""ll_brepnet B-rep face segmentation in native MLX — faithful weight-conversion port. + +This is NOT a simplified re-implementation. It reproduces the EXACT architecture of +``ll_brepnet.models.LLBRepNet`` (the model that reached PyTorch test mIoU 0.828) and +CONVERTS the real trained Lightning checkpoint +(``resources/fusion360/full_train/best.ckpt``) into MLX, so the MLX model *is* the +trained model — same weights, same accuracy — running natively on Apple Silicon. + +Exact architecture (from ll_brepnet.py / uvnet_encoders.py / cadling brep_net.py): + + surface_encoder : 3x [Conv2d(3,pad1) -> BatchNorm2d -> ReLU] -> AdaptiveAvgPool2d(1) (7->32->64->64) + curve_encoder : 3x [Conv1d(3,pad1) -> BatchNorm1d -> ReLU] -> AdaptiveAvgPool1d(1) (6->32->64->64) + face_proj : Linear(8+64 -> 64) -> ReLU edge_proj : Linear(7+64 -> 64) -> ReLU + encoder (BRepNetEncoder): + input_proj : Linear(129 -> 128) -> LayerNorm -> ReLU + 4x layer : residual = h ; h = LayerNorm(relu(W_self·h + W_next·h[next] + + W_prev·h[prev] + W_mate·h[mate])) ; h = h + residual + output_proj: Linear(128 -> 128) (attn_gate feeds only the + coedge->face scatter-mean discarded graph embedding) + seg_head : Linear(128 -> 8) + +Weight conversion details (the non-trivial parts the simplified port got wrong): + * PyTorch Conv2d weight is [out,in,kH,kW] (OIHW); MLX Conv2d is [out,kH,kW,in] (OHWI) -> permute. + * PyTorch Conv1d weight is [out,in,kW] (OIW); MLX Conv1d is [out,kW,in] (OWI) -> permute. + * BatchNorm running_mean / running_var / weight / bias are converted and applied in + inference mode (eps 1e-5) — exactly what the PyTorch model does at eval. + * Linear / LayerNorm map 1:1 (same [out,in] / [dim] layout) with no transpose. + +Both models are driven from the SAME real ``BRepDataset`` (identical z-score +standardization), so the parity comparison is apples-to-apples. + +Modes: probe | convert | parity. +""" + +from __future__ import annotations + +import argparse +import json +import os +import warnings +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import numpy as np # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +from mlx.utils import tree_flatten, tree_unflatten # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] +_DATA = _REPO / "resources/fusion360/full_processed" +_CKPT = _REPO / "resources/fusion360/full_train/best.ckpt" +IGNORE = -1 + + +# --- faithful MLX modules ----------------------------------------------------- +class EvalBN(mlxnn.Module): + """BatchNorm in inference mode over the last (channel) axis, from running stats.""" + + def __init__(self, c, eps=1e-5): + super().__init__() + self.weight = mx.ones((c,)) + self.bias = mx.zeros((c,)) + self.running_mean = mx.zeros((c,)) + self.running_var = mx.ones((c,)) + self.eps = eps + + def __call__(self, x): # x: [..., C] + return (x - self.running_mean) * mx.rsqrt(self.running_var + self.eps) * self.weight + self.bias + + +class SurfEnc(mlxnn.Module): + """UVNetSurfaceEncoder: 3x(Conv2d->BN->ReLU) -> global avg pool. Input [F,U,V,7].""" + + def __init__(self, cin=7, out=64): + super().__init__() + self.c0 = mlxnn.Conv2d(cin, 32, 3, padding=1) + self.b1 = EvalBN(32) + self.c3 = mlxnn.Conv2d(32, 64, 3, padding=1) + self.b4 = EvalBN(64) + self.c6 = mlxnn.Conv2d(64, out, 3, padding=1) + self.b7 = EvalBN(out) + + def __call__(self, g): + x = mlxnn.relu(self.b1(self.c0(g))) + x = mlxnn.relu(self.b4(self.c3(x))) + x = mlxnn.relu(self.b7(self.c6(x))) + return x.mean(axis=(1, 2)) # AdaptiveAvgPool2d(1) -> [F,out] + + +class CurveEnc(mlxnn.Module): + """UVNetCurveEncoder: 3x(Conv1d->BN->ReLU) -> global avg pool. Input [E,U,6].""" + + def __init__(self, cin=6, out=64): + super().__init__() + self.c0 = mlxnn.Conv1d(cin, 32, 3, padding=1) + self.b1 = EvalBN(32) + self.c3 = mlxnn.Conv1d(32, 64, 3, padding=1) + self.b4 = EvalBN(64) + self.c6 = mlxnn.Conv1d(64, out, 3, padding=1) + self.b7 = EvalBN(out) + + def __call__(self, g): + x = mlxnn.relu(self.b1(self.c0(g))) + x = mlxnn.relu(self.b4(self.c3(x))) + x = mlxnn.relu(self.b7(self.c6(x))) + return x.mean(axis=1) # AdaptiveAvgPool1d(1) -> [E,out] + + +class CoedgeConv(mlxnn.Module): + def __init__(self, d): + super().__init__() + self.W_self = mlxnn.Linear(d, d) + self.W_next = mlxnn.Linear(d, d) + self.W_prev = mlxnn.Linear(d, d) + self.W_mate = mlxnn.Linear(d, d) + + def __call__(self, h, nidx, pidx, midx): + return mlxnn.relu(self.W_self(h) + self.W_next(h[nidx]) + self.W_prev(h[pidx]) + self.W_mate(h[midx])) + + +class BRepEncoder(mlxnn.Module): + """cadling BRepNetEncoder: input_proj -> residual coedge convs -> output_proj -> face mean-pool.""" + + def __init__(self, input_dim=129, hidden=128, layers=4): + super().__init__() + self.input_lin = mlxnn.Linear(input_dim, hidden) + self.input_ln = mlxnn.LayerNorm(hidden) + self.conv_layers = [CoedgeConv(hidden) for _ in range(layers)] + self.layer_norms = [mlxnn.LayerNorm(hidden) for _ in range(layers)] + self.output_proj = mlxnn.Linear(hidden, hidden) + + def __call__(self, feats, nidx, pidx, midx, c2f, nf): + h = mlxnn.relu(self.input_ln(self.input_lin(feats))) + for conv, norm in zip(self.conv_layers, self.layer_norms): + res = h + h = norm(conv(h, nidx, pidx, midx)) + res + coedge_emb = self.output_proj(h) + onehot = (c2f[:, None] == mx.arange(nf)[None, :]).astype(coedge_emb.dtype) # [C,F] + return (onehot.T @ coedge_emb) / mx.maximum(onehot.sum(axis=0)[:, None], 1) # [F,hidden] + + +class LLBRepNetMLX(mlxnn.Module): + def __init__(self, nc=8, surf=64, curve=64, ent=64, hidden=128, layers=4): + super().__init__() + self.surface_encoder = SurfEnc(7, surf) + self.curve_encoder = CurveEnc(6, curve) + self.face_proj = mlxnn.Linear(8 + surf, ent) + self.edge_proj = mlxnn.Linear(7 + curve, ent) + self.encoder = BRepEncoder(2 * ent + 1, hidden, layers) + self.seg_head = mlxnn.Linear(hidden, nc) + + def __call__(self, b): + face_repr = mlxnn.relu(self.face_proj( + mx.concatenate([b["ff"], self.surface_encoder(b["fg"])], axis=1))) + edge_repr = mlxnn.relu(self.edge_proj( + mx.concatenate([b["ef"], self.curve_encoder(b["eg"])], axis=1))) + coedge = mx.concatenate([face_repr[b["c2f"]], edge_repr[b["c2e"]], b["rev"]], axis=1) + face_emb = self.encoder(coedge, b["c2n"], b["c2p"], b["c2m"], b["c2f"], b["nf"]) + return self.seg_head(face_emb) + + +# --- weight conversion -------------------------------------------------------- +def convert_checkpoint(ckpt_path, model): + """Load the real Lightning state_dict and assign it into the MLX model.""" + import torch + + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + sd = ck["state_dict"] if "state_dict" in ck else ck + + def lin(key): # Linear / LayerNorm: direct copy + return mx.array(sd[key].detach().cpu().float().numpy()) + + def conv2d(key): # OIHW -> OHWI + return mx.array(np.transpose(sd[key].detach().cpu().float().numpy(), (0, 2, 3, 1))) + + def conv1d(key): # OIW -> OWI + return mx.array(np.transpose(sd[key].detach().cpu().float().numpy(), (0, 2, 1))) + + pairs = [] + # UV-Net encoders: net indices 0(conv) 1(bn) 3(conv) 4(bn) 6(conv) 7(bn) + for enc, convf in (("surface_encoder", conv2d), ("curve_encoder", conv1d)): + for ci, attr in ((0, "c0"), (3, "c3"), (6, "c6")): + pairs.append((f"{enc}.{attr}.weight", convf(f"{enc}.net.{ci}.weight"))) + pairs.append((f"{enc}.{attr}.bias", lin(f"{enc}.net.{ci}.bias"))) + for bi, attr in ((1, "b1"), (4, "b4"), (7, "b7")): + for p in ("weight", "bias", "running_mean", "running_var"): + pairs.append((f"{enc}.{attr}.{p}", lin(f"{enc}.net.{bi}.{p}"))) + # entity projections (Sequential index 0 is the Linear) + pairs += [("face_proj.weight", lin("face_proj.0.weight")), ("face_proj.bias", lin("face_proj.0.bias")), + ("edge_proj.weight", lin("edge_proj.0.weight")), ("edge_proj.bias", lin("edge_proj.0.bias"))] + # coedge encoder + pairs += [("encoder.input_lin.weight", lin("encoder.input_proj.0.weight")), + ("encoder.input_lin.bias", lin("encoder.input_proj.0.bias")), + ("encoder.input_ln.weight", lin("encoder.input_proj.1.weight")), + ("encoder.input_ln.bias", lin("encoder.input_proj.1.bias")), + ("encoder.output_proj.weight", lin("encoder.output_proj.weight")), + ("encoder.output_proj.bias", lin("encoder.output_proj.bias"))] + n_layers = len({k.split("conv_layers.")[1].split(".")[0] + for k in sd if "encoder.conv_layers." in k}) + for i in range(n_layers): + for w in ("W_self", "W_next", "W_prev", "W_mate"): + pairs.append((f"encoder.conv_layers.{i}.{w}.weight", lin(f"encoder.conv_layers.{i}.{w}.weight"))) + pairs.append((f"encoder.conv_layers.{i}.{w}.bias", lin(f"encoder.conv_layers.{i}.{w}.bias"))) + pairs.append((f"encoder.layer_norms.{i}.weight", lin(f"encoder.layer_norms.{i}.weight"))) + pairs.append((f"encoder.layer_norms.{i}.bias", lin(f"encoder.layer_norms.{i}.bias"))) + pairs += [("seg_head.weight", lin("seg_head.weight")), ("seg_head.bias", lin("seg_head.bias"))] + + model.update(tree_unflatten(pairs)) + mx.eval(model.parameters()) + return model, len(pairs) + + +def to_mlx_batch(batch): + fg = np.transpose(batch.face_point_grids.numpy(), (0, 2, 3, 1)) # [F,U,V,7] + eg = np.transpose(batch.edge_point_grids.numpy(), (0, 2, 1)) # [E,U,6] + return { + "ff": mx.array(batch.face_features.numpy().astype(np.float32)), + "ef": mx.array(batch.edge_features.numpy().astype(np.float32)), + "fg": mx.array(fg.astype(np.float32)), "eg": mx.array(eg.astype(np.float32)), + "c2f": mx.array(batch.coedge_to_face.numpy().astype(np.int32)), + "c2e": mx.array(batch.coedge_to_edge.numpy().astype(np.int32)), + "c2n": mx.array(batch.coedge_to_next.numpy().astype(np.int32)), + "c2p": mx.array(batch.coedge_to_prev.numpy().astype(np.int32)), + "c2m": mx.array(batch.coedge_to_mate.numpy().astype(np.int32)), + "rev": mx.array(batch.coedge_reversed.numpy().astype(np.float32)), + "nf": int(batch.face_features.shape[0]), + } + + +def miou(conf, nc): + ious = [] + for c in range(nc): + inter = conf[c, c] + union = conf[c, :].sum() + conf[:, c].sum() - inter + if union > 0: + ious.append(inter / union) + return float(np.mean(ious)) if ious else 0.0 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "convert", "parity"], default="parity") + ap.add_argument("--ckpt", default=str(_CKPT)) + ap.add_argument("--n-test", type=int, default=1500) + ap.add_argument("--out", default=str(_REPO / "ll_brepnet/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + + ds_meta = json.load(open(_DATA / "dataset.json")) + nc = int(ds_meta["num_classes"]) + + model = LLBRepNetMLX(nc=nc) + + if args.mode == "probe": + f = mx.zeros((3, 10, 10, 7)); e = mx.zeros((5, 10, 6)) + b = {"ff": mx.zeros((3, 8)), "ef": mx.zeros((5, 7)), "fg": f, "eg": e, + "c2f": mx.array(np.array([0, 1, 2, 0, 1], np.int32)), + "c2e": mx.array(np.array([0, 1, 2, 3, 4], np.int32)), + "c2n": mx.array(np.array([1, 2, 0, 4, 3], np.int32)), + "c2p": mx.array(np.array([2, 0, 1, 4, 3], np.int32)), + "c2m": mx.array(np.array([3, 4, 2, 0, 1], np.int32)), + "rev": mx.zeros((5, 1)), "nf": 3} + out = model(b) + print(f"probe: logits {out.shape} finite={bool(mx.isfinite(out).all().item())}", flush=True) + return + + model, n = convert_checkpoint(args.ckpt, model) + out_path = f"{args.out}/brepnet_mlx.safetensors" + mx.save_safetensors(out_path, dict(tree_flatten(model.parameters()))) + print(f"converted {n} real tensors -> {out_path}", flush=True) + if args.mode == "convert": + return + + # parity: drive BOTH models from the real BRepDataset + import sys + sys.path.insert(0, str(_REPO / "ll_brepnet")) + sys.path.insert(0, str(_REPO / "cadling")) + import torch + from ll_brepnet.models.ll_brepnet import LLBRepNet + from ll_brepnet.dataloaders.brep_dataset import BRepDataset, brep_collate_fn + + pt = LLBRepNet.load_from_checkpoint(args.ckpt, map_location="cpu") + pt.eval() + ds = BRepDataset(_DATA / "dataset.json", _DATA, split="test", standardize=True) + n_test = min(args.n_test, len(ds)) + print(f"running parity on {n_test} test solids ...", flush=True) + + conf_pt = np.zeros((nc, nc), np.int64) + conf_mlx = np.zeros((nc, nc), np.int64) + agree = total = 0 + for i in range(n_test): + try: + sample = ds[i] + batch = brep_collate_fn([sample]) + except Exception: + continue + with torch.no_grad(): + lp = pt(batch).cpu().numpy() + lm = np.array(model(to_mlx_batch(batch)).tolist()) + lab = batch.labels.numpy() + m = lab != IGNORE + pp, pm = lp.argmax(1), lm.argmax(1) + agree += int((pp[m] == pm[m]).sum()); total += int(m.sum()) + for t, a, b2 in zip(lab[m], pp[m], pm[m]): + conf_pt[t, a] += 1; conf_mlx[t, b2] += 1 + if (i + 1) % 500 == 0: + print(f" {i+1}/{n_test} running mIoU pt={miou(conf_pt,nc):.3f} mlx={miou(conf_mlx,nc):.3f}", flush=True) + + result = {"framework": "MLX (Apple Silicon)", "port": "faithful weight-conversion", + "task": "B-rep face segmentation", "dataset": "Fusion360", "num_classes": nc, + "n_test_solids": total and n_test, "n_faces_scored": int(total), + "source_checkpoint": args.ckpt, + "argmax_agreement_vs_pytorch": round(agree / max(total, 1), 4), + "mlx_mIoU": round(miou(conf_mlx, nc), 4), "pytorch_mIoU": round(miou(conf_pt, nc), 4), + "mlx_acc": round(float(np.trace(conf_mlx) / max(conf_mlx.sum(), 1)), 4), + "pytorch_acc": round(float(np.trace(conf_pt) / max(conf_pt.sum(), 1)), 4), + "pytorch_reference_full_test_mIoU": 0.828, + "checkpoint": out_path} + with open(f"{args.out}/brepnet_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("BREPNET_MLX_PARITY", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_gen/ll_gen/disposal/command_executor.py b/ll_gen/ll_gen/disposal/command_executor.py index 9438c10..7f6fe79 100644 --- a/ll_gen/ll_gen/disposal/command_executor.py +++ b/ll_gen/ll_gen/disposal/command_executor.py @@ -383,83 +383,113 @@ def _execute_sketch_group( def _build_sketch_face( sketch_commands: list[dict[str, Any]], step_offset: int = 0 ) -> TopoDS_Shape | None: - """Build a 2D sketch face from sketch commands. + """Build a CLOSED 2D sketch face from sketch commands. + + Closure-aware construction: a sketch loop is built by THREADING curve + endpoints — each curve starts where the previous one ended — and the loop is + auto-closed by connecting the last endpoint back to the first start. This + guarantees a closed wire even when the upstream generator emits curves whose + absolute endpoints do not coincide. That non-coincidence is the dominant + failure mode for non-autoregressive decoders: each curve's coordinates are + independent argmaxes, so consecutive segments almost never share a vertex and + ``MakeWire`` fails — which is why only self-closing primitives (circles) ever + validated. Threading + auto-closing lets multi-line/arc polygons close, + unlocking diverse valid solids instead of cylinders only. + + A loop consisting solely of circles is self-closing and built directly. Args: sketch_commands: List of sketch commands (LINE, ARC, CIRCLE). step_offset: Offset for step numbering in error messages. Returns: - A TopoDS_Shape representing the sketch face, or None if unsuccessful. - - Raises: - RuntimeError: If edge or wire creation fails. + A TopoDS_Shape (face) for the closed loop, or None if unsuccessful. """ - edges = [] + circle_cmds = [c for c in sketch_commands if c.get("type") == "CIRCLE"] + curve_cmds = [c for c in sketch_commands if c.get("type") in ("LINE", "ARC")] + + # A loop of only circle(s): the circle is itself a closed wire. + if circle_cmds and not curve_cmds: + edge = _create_circle_edge(circle_cmds[0].get("params", {})) + if edge is None: + _log.warning("Failed to create circle edge for sketch loop") + return None + return _face_from_edges([edge]) - for step_idx, cmd in enumerate(sketch_commands): - cmd_type = cmd.get("type", "") - params = cmd.get("params", {}) + if not curve_cmds: + _log.warning("No LINE/ARC/CIRCLE commands in sketch loop") + return None + # Thread endpoints: each curve's start is the previous curve's end. + threaded_edges: list[Any] = [] + current: tuple[float, float] | None = None + first: tuple[float, float] | None = None + for step_idx, cmd in enumerate(curve_cmds): + params = dict(cmd.get("params", {})) + cmd_type = cmd.get("type") try: if cmd_type == "LINE": + start = current if current is not None else (params.get("x1", 0.0), params.get("y1", 0.0)) + end = (params.get("x2", 0.0), params.get("y2", 0.0)) + params.update({"x1": start[0], "y1": start[1], "x2": end[0], "y2": end[1]}) edge = _create_line_edge(params) - elif cmd_type == "ARC": + else: # ARC + start = current if current is not None else (params.get("x_start", 0.0), params.get("y_start", 0.0)) + end = (params.get("x_end", 0.0), params.get("y_end", 0.0)) + params.update({"x_start": start[0], "y_start": start[1], "x_end": end[0], "y_end": end[1]}) edge = _create_arc_edge(params) - elif cmd_type == "CIRCLE": - edge = _create_circle_edge(params) - else: - _log.warning(f"Unknown sketch command type: {cmd_type}") - continue - - if edge is not None: - edges.append(edge) - else: - _log.warning( - f"Failed to create edge for {cmd_type} at step {step_offset + step_idx}" - ) + if edge is None: + # Degenerate arc (collinear/coincident points) -> straight chord. + edge = _create_line_edge({"x1": start[0], "y1": start[1], "x2": end[0], "y2": end[1]}) except Exception as e: - _log.error( - f"Error creating {cmd_type} edge at step {step_offset + step_idx}: {e}" + _log.warning( + f"Skipping {cmd_type} edge at step {step_offset + step_idx}: {e}" ) - raise RuntimeError( - f"Error creating {cmd_type} edge at step {step_offset + step_idx}: {e}" - ) from e - - if not edges: - _log.warning("No valid edges created for sketch") + edge = None + if first is None: + first = start + # Advance the threading point even if this edge was dropped, so the + # remaining curves and the closing segment still connect end-to-end. + current = end + if edge is not None: + threaded_edges.append(edge) + + if not threaded_edges or first is None or current is None: + _log.warning("No valid threaded edges created for sketch loop") return None - # Create wire from edges - try: - wire_maker = BRepBuilderAPI_MakeWire() - for edge in edges: - wire_maker.Add(edge) + # Auto-close: connect the last endpoint back to the first start. + if math.hypot(current[0] - first[0], current[1] - first[1]) > 1e-7: + closing = _create_line_edge( + {"x1": current[0], "y1": current[1], "x2": first[0], "y2": first[1]} + ) + if closing is not None: + threaded_edges.append(closing) - if not wire_maker.IsDone(): - _log.error("Failed to create wire from edges") - return None + return _face_from_edges(threaded_edges) - wire = wire_maker.Wire() - _log.debug(f"Created wire with {len(edges)} edges") - except Exception as e: - _log.error(f"Error creating wire: {e}") - raise RuntimeError(f"Error creating wire: {e}") from e - # Create face from wire - try: - face_maker = BRepBuilderAPI_MakeFace(wire, False) +def _face_from_edges(edges: list[Any]) -> TopoDS_Shape | None: + """Build a wire from connected edges and a planar face from that wire. - if not face_maker.IsDone(): - _log.error("Failed to create face from wire") - return None + Args: + edges: Ordered, end-to-end connected ``TopoDS_Edge`` objects forming a + closed loop. - face = face_maker.Face() - _log.debug("Created face from wire") - return face - except Exception as e: - _log.error(f"Error creating face: {e}") - raise RuntimeError(f"Error creating face: {e}") from e + Returns: + The planar ``TopoDS_Face``, or None if wire/face construction fails. + """ + wire_maker = BRepBuilderAPI_MakeWire() + for edge in edges: + wire_maker.Add(edge) + if not wire_maker.IsDone(): + _log.error("Failed to create wire from edges") + return None + face_maker = BRepBuilderAPI_MakeFace(wire_maker.Wire(), False) + if not face_maker.IsDone(): + _log.error("Failed to create face from wire") + return None + return face_maker.Face() def _create_line_edge(params: dict[str, Any]) -> TopoDS_Shape | None: diff --git a/ll_gen/ll_gen/disposal/surface_executor.py b/ll_gen/ll_gen/disposal/surface_executor.py index 83f0b46..b086c60 100644 --- a/ll_gen/ll_gen/disposal/surface_executor.py +++ b/ll_gen/ll_gen/disposal/surface_executor.py @@ -99,6 +99,26 @@ def execute_latent_proposal(proposal: LatentProposal) -> Any: _log.error(f"Failed to fit B-spline surface for face {i}: {e}") raise + # Primary path: cadling's TopologyMerger is purpose-built to merge + # independently-generated faces (the diffusion case) into a watertight solid + # — it dedups shared edges, averages their geometry, and sews. Use it first; + # fall through to the manual edge-fit/dedup/trim/sew path if it can't close. + if _CADLING_TOPOLOGY_MERGER_AVAILABLE: + try: + merge_res = TopologyMerger().merge(faces) + merged_shape = ( + merge_res.get("shape") if isinstance(merge_res, dict) else None + ) + if merged_shape is not None: + _log.info("Sewed watertight solid via cadling TopologyMerger") + return merged_shape + _log.debug( + "cadling TopologyMerger produced no shape (%s); using manual sew", + merge_res.get("errors") if isinstance(merge_res, dict) else "n/a", + ) + except Exception as e: + _log.debug("cadling TopologyMerger.merge failed (%s); using manual sew", e) + # Step 2: Fit B-spline curves for each edge _log.info("Step 2: Fitting B-spline curves for %d edges", len(proposal.edge_points)) edges = [] @@ -111,15 +131,11 @@ def execute_latent_proposal(proposal: LatentProposal) -> Any: _log.error(f"Failed to fit B-spline curve for edge {i}: {e}") raise - # Step 3: Mating deduplication (topology merger) + # Step 3: Mating deduplication. cadling's TopologyMerger.merge() (tried as the + # primary whole-solid path above) operates on faces, not edges — there is no + # merge_edges(). This manual fallback dedups the fitted edge list directly. _log.info("Step 3: Deduplicating mating edges") - if _CADLING_TOPOLOGY_MERGER_AVAILABLE: - _log.info("Using cadling TopologyMerger for topology merging") - merger = TopologyMerger() - edges = merger.merge_edges(edges) - else: - _log.info("Using built-in edge deduplication") - edges = _deduplicate_edges(edges) + edges = _deduplicate_edges(edges) # Step 4: Surface trimming (edges bound faces) _log.info("Step 4: Trimming surfaces with deduplicated edges") @@ -147,9 +163,25 @@ def _fit_bspline_surface(grid: np.ndarray, tolerance: float = 1e-3) -> Any: RuntimeError: If surface fitting fails. """ if _CADLING_SURFACE_FITTER_AVAILABLE: + # cadling's BSplineSurfaceFitter.fit_surface(point_grid) takes only the + # grid (its tolerance is a constructor arg) and returns a dict + # {'face', 'surface', 'valid', ...} — not a TopoDS_Face. Passing + # tolerance= raised TypeError, so every face silently failed and nothing + # ever sewed. Use the constructor tolerance and extract the face; fall + # through to the direct OCC fit when cadling produces no valid face. _log.debug("Using cadling BSplineSurfaceFitter") - fitter = BSplineSurfaceFitter() - return fitter.fit_surface(grid, tolerance=tolerance) + try: + fitter = BSplineSurfaceFitter(tolerance=tolerance) + except TypeError: + fitter = BSplineSurfaceFitter() + result = fitter.fit_surface(grid) + face = result.get("face") if isinstance(result, dict) else result + if face is not None: + return face + _log.debug( + "cadling fitter produced no face (%s); using direct OCC fit", + result.get("errors") if isinstance(result, dict) else "unknown", + ) if not _OCC_AVAILABLE: raise RuntimeError("pythonocc is required for B-spline surface fitting") diff --git a/ll_gen/ll_gen/training/evaluate_validity.py b/ll_gen/ll_gen/training/evaluate_validity.py index f07c197..5b31dfe 100644 --- a/ll_gen/ll_gen/training/evaluate_validity.py +++ b/ll_gen/ll_gen/training/evaluate_validity.py @@ -250,6 +250,7 @@ def evaluate_validity( output_dir: str | Path = "eval_output", seed: int | None = None, decode_mode: str = "inference", + export: bool = True, ) -> GenerationMetrics: """Measure a generator's disposal-validity rate over a prompt set. @@ -271,6 +272,11 @@ def evaluate_validity( seed: Optional RNG seed for reproducible sampling. decode_mode: ``"inference"`` (``generate``, the deployment path) or ``"training"`` (``generate_for_training``, the path RL optimizes). + export: When True (default), every valid disposed shape is written to + ``/disposed/.step`` and ``.stl`` — so an eval run + produces the generated CAD as real files, not an empty directory. + Pass False for throughput-sensitive callers (e.g. per-step RL evals) + that only need the validity rate. Returns: Populated ``GenerationMetrics`` (``validity_rate`` is the M3 gate). @@ -308,7 +314,7 @@ def evaluate_validity( ) def dispose_fn(proposal: BaseProposal) -> DisposalResult: - return engine.dispose(proposal, export=False) + return engine.dispose(proposal, export=export) model = getattr(generator, "_model", None) results: list[DisposalResult] = [] diff --git a/ll_gen/ll_gen/training/metrics.py b/ll_gen/ll_gen/training/metrics.py index a0f436c..fa62667 100644 --- a/ll_gen/ll_gen/training/metrics.py +++ b/ll_gen/ll_gen/training/metrics.py @@ -103,8 +103,29 @@ def __init__(self, num_bins: int = 64, kernel_bandwidth: float = 0.1) -> None: self.num_bins = num_bins self.kernel_bandwidth = kernel_bandwidth + @staticmethod + def is_valid_solid(result: DisposalResult) -> bool: + """Honest validity gate: a sample counts as valid only if it passes + BRepCheck (``is_valid``) AND forms a non-degenerate solid. + + ``is_valid`` (BRepCheck) alone passes volume-less shells and zero-volume + degenerates, so a generator that emits unsewable faces can score ~1.0 + while producing zero real CAD solids. We therefore additionally require a + closed solid with positive volume whenever a geometry report is present. + Abstract stand-ins without a geometry report (unit tests) fall back to + ``is_valid`` so the rate arithmetic stays testable. + """ + if not getattr(result, "is_valid", False): + return False + gr = getattr(result, "geometry_report", None) + if gr is None: + return True + is_solid = bool(getattr(gr, "is_solid", False) or getattr(gr, "solid_count", 0) >= 1) + vol = getattr(gr, "volume", None) + return is_solid and vol is not None and vol > 1e-4 + def compute_validity_rate(self, results: list[DisposalResult]) -> float: - """Compute fraction of valid samples. + """Compute fraction of valid samples (honest, non-degenerate-solid gated). Args: results: List of disposal results. @@ -114,7 +135,7 @@ def compute_validity_rate(self, results: list[DisposalResult]) -> float: """ if not results: return 0.0 - valid_count = sum(1 for r in results if r.is_valid) + valid_count = sum(1 for r in results if self.is_valid_solid(r)) return valid_count / len(results) def compute_compile_rate(self, results: list[DisposalResult]) -> float: @@ -485,7 +506,7 @@ def compute_all( mean_reward=mean_reward, reward_std=reward_std, num_samples=len(results), - num_valid=sum(1 for r in results if r.is_valid), + num_valid=sum(1 for r in results if self.is_valid_solid(r)), num_compiled=sum(1 for r in results if r.has_shape), num_distinct_valid=num_distinct_valid, ) @@ -509,7 +530,7 @@ def compute_distinct_valid( """ seen: set[Any] = set() for idx, r in enumerate(results): - if not r.is_valid: + if not self.is_valid_solid(r): continue dims = r.geometry_report.bbox_dimensions if r.geometry_report else None if dims is None: diff --git a/ll_gen/mlx/ar_generator_mlx.py b/ll_gen/mlx/ar_generator_mlx.py new file mode 100644 index 0000000..8c6f5f0 --- /dev/null +++ b/ll_gen/mlx/ar_generator_mlx.py @@ -0,0 +1,404 @@ +"""Autoregressive CAD-command generator in native MLX — produces MEASURED-valid CAD. + +The existing ll_gen generators cannot produce valid CAD: the command-VAE's parallel +(z-broadcast) decoder is primitive-limited and posterior-collapses (validity ~0-12%), +and the diffusion path samples faces independently so they never mate (validity 0). +The robust, proven route to valid CAD (DeepCAD / Text2CAD) is to generate the +CONSTRUCTION PROGRAM autoregressively and execute it: the model learns the command +grammar from real data, and the OCC kernel builds the solid command-by-command. + +Pipeline: + real DeepCAD cad_vec -> translate -> command-token sequence (vocab 268: 0=PAD, 1=BOS, + 2=EOS, 6-11=command types, 12..267=quantised param values) + -> causal-transformer LM, teacher-forced next-token training on real sequences + -> autoregressive sampling (temperature + top-k) + -> decode tokens -> command_dicts -> execute_command_proposal -> OCC solid + +Validity is MEASURED through the real kernel and gated HONESTLY against the cylinder +trap: a sample counts as valid only if it forms a solid (solid_count >= 1) with +non-degenerate volume (> eps); we also report num_distinct (rounded bounding boxes) +and the volume spread, so a high rate with one repeated trivial shape is visible. + +Modes: probe | train (train trains, samples, and reports measured validity). +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +import warnings +from collections import Counter +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MPLBACKEND", "Agg") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +matplotlib.use = lambda *a, **k: None + +import numpy as np # noqa: E402 +import h5py # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +import mlx.optimizers as optim # noqa: E402 +from mlx.utils import tree_flatten # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] +_DEEPCAD = str(_REPO / "resources/DeepCAD") +sys.path.insert(0, _DEEPCAD) +sys.path.insert(0, str(_REPO / "resources/ll_gen_proof")) + +# --- tokenization (shared scheme with the stepnet/ocadr trainers) ------------- +LEVELS, RANGE, MAX_LEN = 256, 2.0, 64 +MASK = {"LINE": [0, 1, 2, 3], "ARC": [0, 1, 2, 3, 4, 5], "CIRCLE": [0, 1, 2], + "EXTRUDE": [0, 1, 2, 3, 4, 5, 6, 7], "SOL": [], "EOS": []} +CMD_TOK = {"SOL": 6, "LINE": 7, "ARC": 8, "CIRCLE": 9, "EXTRUDE": 10, "EOS": 11} +TOK_CMD = {v: k for k, v in CMD_TOK.items()} +VOCAB = 12 + LEVELS +PAD, BOS, SEQ_EOS = 0, 1, 2 + + +def _qc(g): + return int(np.clip(round(float(g)), 0, LEVELS - 1)) + + +def _qv(v): + return int(np.clip(round((float(v) + RANGE) / (2 * RANGE) * (LEVELS - 1)), 0, LEVELS - 1)) + + +def _cmds(cad, Circle, Arc): + out = [] + for ext in cad.seq: + for loop in ext.profile.children: + out.append(("SOL", {})) + for cv in loop.children: + if isinstance(cv, Circle): + out.append(("CIRCLE", {0: _qc(cv.center[0]), 1: _qc(cv.center[1]), + 2: _qv(float(cv.radius) / (LEVELS - 1) * 2 * RANGE)})) + elif isinstance(cv, Arc): + s, e, c = cv.start_point, cv.end_point, cv.center + out.append(("ARC", {0: _qc(s[0]), 1: _qc(s[1]), 2: _qc(e[0]), 3: _qc(e[1]), + 4: _qc(c[0]), 5: _qc(c[1])})) + else: + s, e = cv.start_point, cv.end_point + out.append(("LINE", {0: _qc(s[0]), 1: _qc(s[1]), 2: _qc(e[0]), 3: _qc(e[1])})) + out.append(("EXTRUDE", {0: _qv(float(np.clip((abs(float(ext.extent_one)) + + abs(float(ext.extent_two))) * 4, 0.3, 2.0)))})) + out.append(("EOS", {})) + return out + + +def encode_tokens(cmds): + t = [BOS] + for name, slots in cmds: + t.append(CMD_TOK[name]) + for j in MASK[name]: + t.append(12 + int(slots.get(j, 0))) + t.append(SEQ_EOS) + t = t[:MAX_LEN] + return t + [PAD] * (MAX_LEN - len(t)) + + +def decode_tokens(toks): + """Token list -> list of (command_name, {slot: value}). Robust to malformed runs.""" + cmds = [] + i, n = 0, len(toks) + while i < n: + t = int(toks[i]) + if t == SEQ_EOS or t == PAD: + break + if t in TOK_CMD: + name = TOK_CMD[t] + if name == "EOS": + break + slots = {} + ok = True + for j in MASK[name]: + i += 1 + if i < n and 12 <= int(toks[i]) < 12 + LEVELS: + slots[j] = int(toks[i]) - 12 + else: + ok = False + break + if ok: + cmds.append((name, slots)) + i += 1 + else: + i += 1 # stray param token without a command — skip + return cmds + + +def command_dicts(cmds): + out = [] + for name, slots in cmds: + p = [0] * 16 + m = [False] * 16 + for j in MASK[name]: + p[j] = int(slots.get(j, 0)) + m[j] = True + out.append({"command_type": name, "parameters": p, "parameter_mask": m}) + return out + + +def build_dataset(n_target, cache): + if cache and os.path.exists(cache): + d = np.load(cache) + if d["tokens"].shape[0] >= n_target: + return d["tokens"][:n_target] + from cadlib.extrude import CADSequence + from cadlib.curves import Arc, Circle + + toks = [] + for f in sorted(glob.glob(os.path.join(_DEEPCAD, "data/cad_vec/*/*.h5"))): + if len(toks) >= n_target: + break + try: + with h5py.File(f, "r") as h: + vec = h["vec"][:].astype(int) + cad = CADSequence.from_vector(vec, is_numerical=True, n=256) + cmds = _cmds(cad, Circle, Arc) + enc = encode_tokens(cmds) + if enc[1] != PAD: # non-empty + toks.append(enc) + if len(toks) % 5000 == 0 and len(toks): + print(f" built {len(toks)}/{n_target}", flush=True) + except Exception: + continue + toks = np.array(toks, np.int32) + if cache: + np.savez(cache, tokens=toks) + return toks + + +# --- MLX causal-transformer language model ------------------------------------ +class CausalBlock(mlxnn.Module): + def __init__(self, d, heads, ff): + super().__init__() + self.h, self.hd = heads, d // heads + self.qkv = mlxnn.Linear(d, 3 * d) + self.proj = mlxnn.Linear(d, d) + self.n1 = mlxnn.LayerNorm(d) + self.n2 = mlxnn.LayerNorm(d) + self.fc1 = mlxnn.Linear(d, ff) + self.fc2 = mlxnn.Linear(ff, d) + self.d = d + + def __call__(self, x, mask): + b, s, _ = x.shape + h = self.n1(x) + q, k, v = mx.split(self.qkv(h), 3, axis=-1) + + def sp(t): + return mx.transpose(t.reshape(b, s, self.h, self.hd), (0, 2, 1, 3)) + + q, k, v = sp(q), sp(k), sp(v) + att = (q @ mx.transpose(k, (0, 1, 3, 2))) / (self.hd ** 0.5) + mask + ctx = mx.softmax(att, axis=-1) @ v + ctx = mx.transpose(ctx, (0, 2, 1, 3)).reshape(b, s, self.d) + x = x + self.proj(ctx) + return x + self.fc2(mlxnn.gelu(self.fc1(self.n2(x)))) + + +class ARGPT(mlxnn.Module): + def __init__(self, vocab=VOCAB, d=256, layers=6, heads=8, ff=1024, maxlen=MAX_LEN): + super().__init__() + self.embed = mlxnn.Embedding(vocab, d) + self.pos = mx.zeros((1, maxlen, d)) + self.blocks = [CausalBlock(d, heads, ff) for _ in range(layers)] + self.norm = mlxnn.LayerNorm(d) + self.head = mlxnn.Linear(d, vocab) + self.maxlen = maxlen + + def __call__(self, ids): + s = ids.shape[1] + mask = mx.where(mx.triu(mx.ones((s, s)), k=1) > 0, + mx.array(-1e9, mx.float32), mx.array(0.0, mx.float32))[None, None] + x = self.embed(ids) + self.pos[:, :s, :] + for blk in self.blocks: + x = blk(x, mask) + return self.head(self.norm(x)) + + +def sample(model, n, temperature=1.0, top_k=20): + """Autoregressive batch sampling -> list of token lists (BOS-stripped).""" + cur = mx.full((n, 1), BOS, dtype=mx.int32) + done = np.zeros(n, bool) + seqs = [[] for _ in range(n)] + for _ in range(MAX_LEN - 1): + logits = model(cur)[:, -1, :] / temperature # [n, vocab] + if top_k: + kth = mx.sort(logits, axis=-1)[:, -top_k][:, None] + logits = mx.where(logits < kth, mx.array(-1e9, mx.float32), logits) + nxt = mx.random.categorical(logits) # [n] + mx.eval(nxt) + nxt_np = np.array(nxt.tolist()) + for i in range(n): + if not done[i]: + t = int(nxt_np[i]) + seqs[i].append(t) + if t == SEQ_EOS or t == CMD_TOK["EOS"]: + done[i] = True + cur = mx.concatenate([cur, nxt[:, None].astype(mx.int32)], axis=1) + if done.all(): + break + return seqs + + +# --- honest validity through the real OCC kernel ------------------------------ +def make_evaluator(): + from ll_gen.proposals.command_proposal import CommandSequenceProposal + from ll_gen.disposal.command_executor import execute_command_proposal + from OCC.Core.GProp import GProp_GProps + from OCC.Core.BRepGProp import brepgprop + from OCC.Core.TopExp import TopExp_Explorer + from OCC.Core.TopAbs import TopAbs_SOLID + from OCC.Core.Bnd import Bnd_Box + from OCC.Core.BRepBndLib import brepbndlib + + def evaluate(toks): + """Return (is_valid_solid, volume, bbox_signature) for one token sequence.""" + cmds = decode_tokens(toks) + if not cmds: + return False, 0.0, None + try: + shape = execute_command_proposal(CommandSequenceProposal( + command_dicts=command_dicts(cmds), quantization_bits=8, normalization_range=2.0)) + except Exception: + return False, 0.0, None + if shape is None: + return False, 0.0, None + nsolids = 0 + e = TopExp_Explorer(shape, TopAbs_SOLID) + while e.More(): + nsolids += 1 + e.Next() + if nsolids < 1: + return False, 0.0, None + props = GProp_GProps() + brepgprop.VolumeProperties(shape, props) + vol = abs(props.Mass()) + if vol <= 1e-4: # reject zero-volume degenerates + return False, vol, None + box = Bnd_Box() + brepbndlib.Add(shape, box) + xmin, ymin, zmin, xmax, ymax, zmax = box.Get() + sig = (round(xmax - xmin, 1), round(ymax - ymin, 1), round(zmax - zmin, 1)) + return True, vol, sig + + return evaluate + + +def measure_validity(model, evaluate, n, temperature, top_k): + seqs = sample(model, n, temperature, top_k) + valid, vols, sigs = 0, [], [] + for s in seqs: + ok, vol, sig = evaluate(s) + if ok: + valid += 1 + vols.append(vol) + sigs.append(sig) + distinct = len(set(sigs)) + # cylinder-trap guard: fraction of valid samples that are the single most common shape + top_frac = (Counter(sigs).most_common(1)[0][1] / len(sigs)) if sigs else 0.0 + return {"n": n, "validity": valid / n, "num_valid": valid, "num_distinct": distinct, + "top_shape_frac": round(top_frac, 3), + "mean_volume": float(np.mean(vols)) if vols else 0.0, + "vol_p10_p90": [float(np.percentile(vols, 10)), float(np.percentile(vols, 90))] if vols else [0, 0]} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "train"], default="train") + ap.add_argument("--n-train", type=int, default=40000) + ap.add_argument("--epochs", type=int, default=40) + ap.add_argument("--bs", type=int, default=128) + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--n-eval", type=int, default=128) + ap.add_argument("--temperature", type=float, default=1.0) + ap.add_argument("--top-k", type=int, default=20) + ap.add_argument("--out", default=str(_REPO / "ll_gen/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + cache = f"{args.out}/ar_tokens_cache.npz" + + if args.mode == "probe": + toks = build_dataset(64, None) + model = ARGPT() + out = model(mx.array(toks[:4])) + ev = make_evaluator() + seqs = sample(model, 4) + v0 = [ev(s)[0] for s in seqs] + print(f"probe: data {toks.shape}, logits {out.shape}, sampled-valid(untrained)={sum(v0)}/4", flush=True) + return + + print("building/loading real DeepCAD command sequences ...", flush=True) + toks = build_dataset(args.n_train, cache) + print(f"dataset: {toks.shape[0]} sequences", flush=True) + n_val = min(2000, toks.shape[0] // 10) + tr = toks[n_val:] + model = ARGPT() + opt = optim.AdamW(learning_rate=args.lr, weight_decay=0.01) + + def loss_fn(ids): + logits = model(ids[:, :-1]) + tgt = ids[:, 1:] + mask = (tgt != PAD).astype(mx.float32) + ce = mlxnn.losses.cross_entropy(logits.reshape(-1, VOCAB), tgt.reshape(-1), reduction="none") + ce = ce.reshape(tgt.shape) * mask + return ce.sum() / mx.maximum(mask.sum(), 1) + + lg = mlxnn.value_and_grad(model, loss_fn) + evaluate = make_evaluator() + print("measuring untrained baseline validity ...", flush=True) + base = measure_validity(model, evaluate, args.n_eval, args.temperature, args.top_k) + print(f"BASELINE (untrained): {json.dumps(base)}", flush=True) + + n = tr.shape[0] + best = -1.0 + for epoch in range(args.epochs): + perm = np.random.permutation(n) + tot = 0.0 + nb = 0 + for k in range(0, n, args.bs): + idx = perm[k:k + args.bs] + lv, g = lg(mx.array(tr[idx])) + opt.update(model, g) + mx.eval(model.parameters(), opt.state, lv) + tot += float(lv.item()) + nb += 1 + if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1: + m = measure_validity(model, evaluate, args.n_eval, args.temperature, args.top_k) + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} " + f"validity={m['validity']:.3f} valid={m['num_valid']} distinct={m['num_distinct']} " + f"mean_vol={m['mean_volume']:.3f}", flush=True) + if m["validity"] > best: + best = m["validity"] + mx.save_safetensors(f"{args.out}/ar_generator_mlx.safetensors", + dict(tree_flatten(model.parameters()))) + else: + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f}", flush=True) + + final = measure_validity(model, evaluate, max(args.n_eval, 256), args.temperature, args.top_k) + result = {"framework": "MLX", "model": "autoregressive CAD-command transformer (DeepCAD-style)", + "task": "generate valid CAD construction programs", "dataset": "DeepCAD cad_vec", + "n_train": int(tr.shape[0]), "epochs": args.epochs, "vocab": VOCAB, + "sampling": {"temperature": args.temperature, "top_k": args.top_k}, + "validity_gate": "is_solid AND volume>1e-4 (non-degenerate), measured via real OCC kernel", + "baseline_untrained_validity": round(base["validity"], 4), + "best_validity": round(best, 4), "final": final, + "checkpoint": f"{args.out}/ar_generator_mlx.safetensors"} + with open(f"{args.out}/ar_generator_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("AR_GENERATOR_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_gen/mlx/latent_diffusion_mlx.py b/ll_gen/mlx/latent_diffusion_mlx.py new file mode 100644 index 0000000..23c90d5 --- /dev/null +++ b/ll_gen/mlx/latent_diffusion_mlx.py @@ -0,0 +1,359 @@ +"""Latent diffusion that produces VALID CAD — the real fix for the diffusion path. + +The shipped diffusion (ll_stepnet StructuredDiffusion + GeometryCodec) denoises raw +B-rep geometry — independent face UV-grids + edge polylines — then tries to SEW them. +Independently-sampled faces never share exact boundaries, so the sewer cannot close a +solid: honest validity 0.0 (its own docstring documents the dead-end). + +The robust fix (DeepCAD's actual generative design) changes the REPRESENTATION: diffuse +in the latent of a CAD-PROGRAM autoencoder, and decode with an execution-respecting +autoregressive decoder. The decoder emits a construction program the OCC kernel builds +into a watertight solid, so validity is high — and it comes from the decoder; the +diffusion supplies the latent prior (controllable, unconditional sampling). + +Architecture: + SeqAutoencoder (deterministic, NOT a VAE — avoids posterior collapse): + encoder : embed + bidirectional transformer + mean-pool -> Linear -> z [d_z] + decoder : z-conditioned causal transformer (z added at every position); trained + teacher-forced with WORD-DROPOUT on the decoder inputs so the latent must + carry the global program (else the AR decoder ignores z). + LatentDDPM : a denoiser MLP over the (normalised) z's; standard DDPM eps-prediction. + +Honest acceptance bar (set in advance): validity of DIFFUSION-SAMPLED z +(z_T -> denoise -> z_0 -> decode -> execute), through the real OCC kernel gated on a +non-degenerate solid (solid + volume>1e-4), with num_distinct > 1, beating the z=0 +predict-the-mean baseline. Reconstruction validity is reported too but is NOT the bar. + +Modes: probe | train (train: AE -> DDPM -> measured sampled-z validity). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np +import mlx.core as mx +import mlx.nn as mlxnn +import mlx.optimizers as optim +from mlx.utils import tree_flatten + +_REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(_REPO / "ll_gen/mlx")) + +# reuse the AR generator's validated tokenizer / executor / honest validity gate +from ar_generator_mlx import ( # noqa: E402 + MAX_LEN, VOCAB, PAD, BOS, SEQ_EOS, CMD_TOK, + build_dataset, decode_tokens, command_dicts, make_evaluator, +) +from collections import Counter # noqa: E402 + +MASK_TOK = 3 # unused id in the vocab, used as the word-dropout placeholder + + +# --- transformer blocks ------------------------------------------------------- +class Block(mlxnn.Module): + def __init__(self, d, heads, ff): + super().__init__() + self.h, self.hd, self.d = heads, d // heads, d + self.qkv = mlxnn.Linear(d, 3 * d) + self.proj = mlxnn.Linear(d, d) + self.n1 = mlxnn.LayerNorm(d) + self.n2 = mlxnn.LayerNorm(d) + self.fc1 = mlxnn.Linear(d, ff) + self.fc2 = mlxnn.Linear(ff, d) + + def __call__(self, x, causal): + b, s, _ = x.shape + h = self.n1(x) + q, k, v = mx.split(self.qkv(h), 3, axis=-1) + + def sp(t): + return mx.transpose(t.reshape(b, s, self.h, self.hd), (0, 2, 1, 3)) + + q, k, v = sp(q), sp(k), sp(v) + att = (q @ mx.transpose(k, (0, 1, 3, 2))) / (self.hd ** 0.5) + if causal is not None: + att = att + causal + ctx = mx.softmax(att, axis=-1) @ v + ctx = mx.transpose(ctx, (0, 2, 1, 3)).reshape(b, s, self.d) + x = x + self.proj(ctx) + return x + self.fc2(mlxnn.gelu(self.fc1(self.n2(x)))) + + +class SeqAutoencoder(mlxnn.Module): + def __init__(self, vocab=VOCAB, d=256, d_z=64, enc_layers=3, dec_layers=4, heads=8, ff=1024): + super().__init__() + self.embed = mlxnn.Embedding(vocab, d) + self.pos = mx.zeros((1, MAX_LEN, d)) + self.enc_blocks = [Block(d, heads, ff) for _ in range(enc_layers)] + self.enc_norm = mlxnn.LayerNorm(d) + self.to_z = mlxnn.Linear(d, d_z) + self.from_z = mlxnn.Linear(d_z, d) + self.dec_blocks = [Block(d, heads, ff) for _ in range(dec_layers)] + self.dec_norm = mlxnn.LayerNorm(d) + self.head = mlxnn.Linear(d, vocab) + self.d, self.d_z = d, d_z + + def encode(self, ids): + x = self.embed(ids) + self.pos[:, : ids.shape[1], :] + for blk in self.enc_blocks: + x = blk(x, None) # bidirectional + x = self.enc_norm(x) + m = (ids != PAD).astype(x.dtype)[..., None] + pooled = (x * m).sum(axis=1) / mx.maximum(m.sum(axis=1), 1) + return self.to_z(pooled) # [B, d_z] + + def decode(self, ids, z): + s = ids.shape[1] + causal = mx.where(mx.triu(mx.ones((s, s)), k=1) > 0, + mx.array(-1e9, mx.float32), mx.array(0.0, mx.float32))[None, None] + zc = self.from_z(z)[:, None, :] # [B,1,d] broadcast over positions + x = self.embed(ids) + self.pos[:, :s, :] + zc + for blk in self.dec_blocks: + x = blk(x, causal) + return self.head(self.dec_norm(x)) + + +# --- latent DDPM -------------------------------------------------------------- +class TimeEmbed(mlxnn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + self.fc1 = mlxnn.Linear(dim, dim) + self.fc2 = mlxnn.Linear(dim, dim) + + def __call__(self, t): # t: [B] in [0,1] + half = self.dim // 2 + freqs = mx.exp(-np.log(10000) * mx.arange(half) / half) + a = t[:, None] * freqs[None] + emb = mx.concatenate([mx.sin(a), mx.cos(a)], axis=-1) + return self.fc2(mlxnn.gelu(self.fc1(emb))) + + +class Denoiser(mlxnn.Module): + def __init__(self, d_z=64, hidden=512, tdim=128): + super().__init__() + self.time = TimeEmbed(tdim) + self.inp = mlxnn.Linear(d_z + tdim, hidden) + self.h1 = mlxnn.Linear(hidden, hidden) + self.h2 = mlxnn.Linear(hidden, hidden) + self.out = mlxnn.Linear(hidden, d_z) + + def __call__(self, z, t): + te = self.time(t) + x = mlxnn.gelu(self.inp(mx.concatenate([z, te], axis=-1))) + x = x + mlxnn.gelu(self.h1(x)) + x = x + mlxnn.gelu(self.h2(x)) + return self.out(x) + + +class DDPM: + """Standard DDPM schedule + sampling for a flat latent vector.""" + + def __init__(self, steps=200): + self.steps = steps + betas = np.linspace(1e-4, 0.02, steps).astype(np.float32) + alphas = 1.0 - betas + abar = np.cumprod(alphas) + self.betas = mx.array(betas) + self.alphas = mx.array(alphas) + self.abar = mx.array(abar) + self._abar_np = abar + + def q_sample(self, z0, t_idx, noise): + ab = mx.sqrt(self.abar[t_idx])[:, None] + omab = mx.sqrt(1 - self.abar[t_idx])[:, None] + return ab * z0 + omab * noise + + def sample(self, denoiser, n, d_z): + z = mx.random.normal((n, d_z)) + for i in range(self.steps - 1, -1, -1): + t = mx.full((n,), i / self.steps) + eps = denoiser(z, t) + a = self.alphas[i] + ab = self.abar[i] + coef = (1 - a) / mx.sqrt(1 - ab) + mean = (z - coef * eps) / mx.sqrt(a) + if i > 0: + z = mean + mx.sqrt(self.betas[i]) * mx.random.normal((n, d_z)) + else: + z = mean + mx.eval(z) + return z + + +# --- decoding (autoregressive, conditioned on z) ------------------------------ +def ar_decode(ae, z, temperature=1.0, top_k=20): + n = z.shape[0] + cur = mx.full((n, 1), BOS, dtype=mx.int32) + done = np.zeros(n, bool) + seqs = [[] for _ in range(n)] + for _ in range(MAX_LEN - 1): + logits = ae.decode(cur, z)[:, -1, :] / temperature + if top_k: + kth = mx.sort(logits, axis=-1)[:, -top_k][:, None] + logits = mx.where(logits < kth, mx.array(-1e9, mx.float32), logits) + nxt = mx.random.categorical(logits) + mx.eval(nxt) + nn = np.array(nxt.tolist()) + for i in range(n): + if not done[i]: + t = int(nn[i]) + seqs[i].append(t) + if t == SEQ_EOS or t == CMD_TOK["EOS"]: + done[i] = True + cur = mx.concatenate([cur, nxt[:, None].astype(mx.int32)], axis=1) + if done.all(): + break + return seqs + + +def measure(seqs, evaluate): + valid, vols, sigs = 0, [], [] + for s in seqs: + ok, vol, sig = evaluate(s) + if ok: + valid += 1 + vols.append(vol) + sigs.append(sig) + top = (Counter(sigs).most_common(1)[0][1] / len(sigs)) if sigs else 0.0 + return {"validity": valid / len(seqs), "num_valid": valid, "num_distinct": len(set(sigs)), + "top_shape_frac": round(top, 3), "mean_volume": float(np.mean(vols)) if vols else 0.0} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "train"], default="train") + ap.add_argument("--n-train", type=int, default=40000) + ap.add_argument("--ae-epochs", type=int, default=40) + ap.add_argument("--dm-epochs", type=int, default=300) + ap.add_argument("--bs", type=int, default=128) + ap.add_argument("--d-z", type=int, default=64) + ap.add_argument("--word-dropout", type=float, default=0.5) + ap.add_argument("--ddpm-steps", type=int, default=200) + ap.add_argument("--n-eval", type=int, default=256) + ap.add_argument("--out", default=str(_REPO / "ll_gen/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + cache = f"{args.out}/ar_tokens_cache.npz" + + if args.mode == "probe": + toks = build_dataset(256, cache) + ae = SeqAutoencoder(d_z=args.d_z) + z = ae.encode(mx.array(toks[:8])) + logits = ae.decode(mx.array(toks[:8]), z) + print(f"probe: tokens {toks.shape} z {z.shape} dec_logits {logits.shape}", flush=True) + return + + print("loading real DeepCAD command sequences ...", flush=True) + toks = build_dataset(args.n_train, cache) + print(f"{toks.shape[0]} sequences", flush=True) + evaluate = make_evaluator() + ae = SeqAutoencoder(d_z=args.d_z) + opt = optim.AdamW(learning_rate=3e-4, weight_decay=0.01) + + def ae_loss(ids): + # word-dropout on decoder inputs forces the latent to carry global structure + z = ae.encode(ids) + din = ids[:, :-1] + keep = (mx.random.uniform(shape=din.shape) > args.word_dropout) + din = mx.where(keep, din, mx.array(MASK_TOK, mx.int32)) + logits = ae.decode(din, z) + tgt = ids[:, 1:] + m = (tgt != PAD).astype(mx.float32) + ce = mlxnn.losses.cross_entropy(logits.reshape(-1, VOCAB), tgt.reshape(-1), reduction="none") + ce = ce.reshape(tgt.shape) * m + zreg = 1e-4 * (z * z).mean() # keep the latent bounded -> diffusable + return ce.sum() / mx.maximum(m.sum(), 1) + zreg + + lg = mlxnn.value_and_grad(ae, ae_loss) + n = toks.shape[0] + print("training the program autoencoder ...", flush=True) + for epoch in range(args.ae_epochs): + perm = np.random.permutation(n) + tot = 0.0 + nb = 0 + for k in range(0, n, args.bs): + idx = perm[k:k + args.bs] + lv, g = lg(mx.array(toks[idx])) + opt.update(ae, g) + mx.eval(ae.parameters(), opt.state, lv) + tot += float(lv.item()) + nb += 1 + if (epoch + 1) % 10 == 0 or epoch == args.ae_epochs - 1: + # reconstruction validity (greedy, teacher-free) — sanity, NOT the bar + zr = ae.encode(mx.array(toks[:args.n_eval])) + rec = measure(ar_decode(ae, zr, temperature=0.7), evaluate) + print(f"AE epoch {epoch+1}/{args.ae_epochs} loss={tot/max(nb,1):.4f} " + f"recon_validity={rec['validity']:.3f} distinct={rec['num_distinct']}", flush=True) + + # encode the corpus, normalise the latent for diffusion + print("encoding corpus -> latent bank ...", flush=True) + zs = [] + for k in range(0, n, 512): + zs.append(np.array(ae.encode(mx.array(toks[k:k + 512])).tolist())) + zbank = np.concatenate(zs, axis=0).astype(np.float32) + zmean = zbank.mean(0, keepdims=True) + zstd = zbank.std(0, keepdims=True) + 1e-6 + zn = (zbank - zmean) / zstd + zmean_mx, zstd_mx = mx.array(zmean), mx.array(zstd) + + # train the latent DDPM + print("training the latent DDPM ...", flush=True) + den = Denoiser(d_z=args.d_z) + ddpm = DDPM(steps=args.ddpm_steps) + dopt = optim.AdamW(learning_rate=3e-4, weight_decay=0.0) + + def dm_loss(z0): + b = z0.shape[0] + t_idx = mx.array(np.random.randint(0, args.ddpm_steps, b)) + noise = mx.random.normal(z0.shape) + zt = ddpm.q_sample(z0, t_idx, noise) + eps = den(zt, t_idx.astype(mx.float32) / args.ddpm_steps) + return ((eps - noise) ** 2).mean() + + dlg = mlxnn.value_and_grad(den, dm_loss) + m = zn.shape[0] + for epoch in range(args.dm_epochs): + perm = np.random.permutation(m) + tot = 0.0 + nb = 0 + for k in range(0, m, args.bs): + lv, g = dlg(mx.array(zn[perm[k:k + args.bs]])) + dopt.update(den, g) + mx.eval(den.parameters(), dopt.state, lv) + tot += float(lv.item()) + nb += 1 + if (epoch + 1) % 50 == 0 or epoch == args.dm_epochs - 1: + print(f"DDPM epoch {epoch+1}/{args.dm_epochs} loss={tot/max(nb,1):.5f}", flush=True) + + # === the acceptance bar: validity of DIFFUSION-SAMPLED z === + print("sampling z ~ diffusion -> decode -> execute ...", flush=True) + zsamp = ddpm.sample(den, args.n_eval, args.d_z) * zstd_mx + zmean_mx + sampled = measure(ar_decode(ae, zsamp, temperature=1.0), evaluate) + # predict-the-mean baseline: z = 0 (normalised) -> unnormalise -> decode + zzero = mx.broadcast_to(zmean_mx, (args.n_eval, args.d_z)) + zerob = measure(ar_decode(ae, zzero, temperature=1.0), evaluate) + + result = {"framework": "MLX", "model": "latent DDPM over a z-conditioned AR program autoencoder", + "fix": "diffuse the construction-program latent + execution-respecting AR decoder " + "(replaces independent-face geometry diffusion that could not sew)", + "n_train": int(n), "d_z": args.d_z, "ae_epochs": args.ae_epochs, + "dm_epochs": args.dm_epochs, "ddpm_steps": args.ddpm_steps, + "validity_gate": "is_solid AND volume>1e-4, real OCC kernel", + "sampled_z_validity": round(sampled["validity"], 4), "sampled_z": sampled, + "z0_mean_baseline_validity": round(zerob["validity"], 4), "z0_baseline": zerob, + "checkpoint": f"{args.out}/latent_diffusion_mlx.safetensors"} + mx.save_safetensors(f"{args.out}/latent_diffusion_mlx.safetensors", + dict(tree_flatten(ae.parameters())) | {f"den.{k}": v for k, v in tree_flatten(den.parameters())}) + with open(f"{args.out}/latent_diffusion_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("LATENT_DIFFUSION_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_gen/mlx/vae_mlx.py b/ll_gen/mlx/vae_mlx.py new file mode 100644 index 0000000..b3f66f4 --- /dev/null +++ b/ll_gen/mlx/vae_mlx.py @@ -0,0 +1,301 @@ +"""ll_gen STEPVAE in native MLX — faithful weight-conversion port. + +The ll_gen neural generator is ``ll_stepnet.stepnet.vae.STEPVAE`` (a transformer +encoder-decoder VAE over CAD command-token sequences), trained to the checkpoints +``ll_gen/checkpoints/vae_warm.pt`` / ``vae_rl_solid.pt``. This reproduces that EXACT +architecture in MLX and CONVERTS the real trained weights, so the MLX VAE *is* the +trained model — same weights, same outputs — running natively on Apple Silicon. + +Architecture (from the checkpoint + vae.py): + encoder : STEPTransformerEncoder (token_emb 50000x256, pos 5000x256, 6 post-norm + TransformerEncoderLayers, final LayerNorm) -> masked-mean pool + mu_head / log_var_head : Linear(256->256) (log_var clamped to [-30, 20]) + decode(z): + z_proj = latent_project(z) # Linear 256->256 + hidden = z_proj.broadcast[B,S,256] + dec_pos_embedding[1,60,256] + z_memory = z_proj.broadcast[B,S,256] + hidden = decoder._transformer(hidden, memory=z_memory, tgt_mask=causal) # 6x + TransformerDecoderLayer: causal self_attn + cross multihead_attn + FFN, + post-norm with 3 LayerNorms + hidden = decoder.layer_norm(hidden) + command_head : Linear(256->6) param_heads : 16 x Linear(256->256) + +Notes: + * decode is PARALLEL — the decoder input is z broadcast over positions, NOT shifted + tokens, so decoder.token_embedding / pos_embedding are UNUSED on this path (skipped). + * eval reparam returns mu (deterministic), so encode->decode(mu) is reproducible. + * MHA is implemented manually and splits the packed in_proj_weight into Wq/Wk/Wv so + the same module serves self-attention AND cross-attention (q=tgt, k=v=memory). + +Modes: probe | convert | parity. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import warnings +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import numpy as np # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +from mlx.utils import tree_flatten, tree_unflatten # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] + + +# --- MLX modules -------------------------------------------------------------- +class MHA(mlxnn.Module): + """nn.MultiheadAttention from a packed in_proj [3d,d]; splits Wq/Wk/Wv so it + serves both self-attention and cross-attention. Optional additive mask.""" + + def __init__(self, d, heads): + super().__init__() + self.d, self.h, self.hd = d, heads, d // heads + self.in_proj = mlxnn.Linear(d, 3 * d) + self.out_proj = mlxnn.Linear(d, d) + + def __call__(self, query, key, mask=None): + d = self.d + w, bvec = self.in_proj.weight, self.in_proj.bias + wq, wk, wv = w[:d], w[d:2 * d], w[2 * d:] + bq, bk, bv = bvec[:d], bvec[d:2 * d], bvec[2 * d:] + q = query @ wq.T + bq + k = key @ wk.T + bk + v = key @ wv.T + bv + b, s, _ = q.shape + sk = k.shape[1] + + def split(t, n): + return mx.transpose(t.reshape(b, n, self.h, self.hd), (0, 2, 1, 3)) + + q, k, v = split(q, s), split(k, sk), split(v, sk) + scores = (q @ mx.transpose(k, (0, 1, 3, 2))) / (self.hd ** 0.5) # [b,h,s,sk] + if mask is not None: + scores = scores + mask + ctx = mx.softmax(scores, axis=-1) @ v + ctx = mx.transpose(ctx, (0, 2, 1, 3)).reshape(b, s, d) + return self.out_proj(ctx) + + +class EncLayer(mlxnn.Module): + """nn.TransformerEncoderLayer, post-norm, relu.""" + + def __init__(self, d=256, heads=8, ff=1024): + super().__init__() + self.self_attn = MHA(d, heads) + self.linear1 = mlxnn.Linear(d, ff) + self.linear2 = mlxnn.Linear(ff, d) + self.norm1 = mlxnn.LayerNorm(d) + self.norm2 = mlxnn.LayerNorm(d) + + def __call__(self, x): + x = self.norm1(x + self.self_attn(x, x)) + return self.norm2(x + self.linear2(mlxnn.relu(self.linear1(x)))) + + +class EncoderMLX(mlxnn.Module): + def __init__(self, vocab=50000, d=256, layers=6, heads=8, ff=1024): + super().__init__() + self.token_embedding = mlxnn.Embedding(vocab, d) + self.pos_embedding = mx.zeros((1, 5000, d)) + self.layers = [EncLayer(d, heads, ff) for _ in range(layers)] + self.layer_norm = mlxnn.LayerNorm(d) + + def __call__(self, ids): + x = self.token_embedding(ids) + self.pos_embedding[:, : ids.shape[1], :] + for layer in self.layers: + x = layer(x) + return self.layer_norm(x) + + +class DecLayer(mlxnn.Module): + """nn.TransformerDecoderLayer, post-norm, relu: causal self-attn + cross-attn + FFN.""" + + def __init__(self, d=256, heads=8, ff=1024): + super().__init__() + self.self_attn = MHA(d, heads) + self.multihead_attn = MHA(d, heads) + self.linear1 = mlxnn.Linear(d, ff) + self.linear2 = mlxnn.Linear(ff, d) + self.norm1 = mlxnn.LayerNorm(d) + self.norm2 = mlxnn.LayerNorm(d) + self.norm3 = mlxnn.LayerNorm(d) + + def __call__(self, tgt, memory, causal): + tgt = self.norm1(tgt + self.self_attn(tgt, tgt, mask=causal)) + tgt = self.norm2(tgt + self.multihead_attn(tgt, memory)) + return self.norm3(tgt + self.linear2(mlxnn.relu(self.linear1(tgt)))) + + +class STEPVAEMLX(mlxnn.Module): + def __init__(self, latent=256, d=256, layers=6, heads=8, ff=1024, + max_seq=60, n_cmd=6, n_param=16, n_levels=256): + super().__init__() + self.encoder = EncoderMLX(50000, d, layers, heads, ff) + self.mu_head = mlxnn.Linear(d, latent) + self.log_var_head = mlxnn.Linear(d, latent) + self.latent_project = mlxnn.Linear(latent, d) + self.dec_pos_embedding = mx.zeros((1, max_seq, d)) + self.dec_layers = [DecLayer(d, heads, ff) for _ in range(layers)] + self.dec_layer_norm = mlxnn.LayerNorm(d) + self.command_head = mlxnn.Linear(d, n_cmd) + self.param_heads = [mlxnn.Linear(d, n_levels) for _ in range(n_param)] + self.d = d + + def encode(self, ids): + hidden = self.encoder(ids) + pooled = hidden.mean(axis=1) # no padding mask -> plain mean + mu = self.mu_head(pooled) + log_var = mx.clip(self.log_var_head(pooled), -30, 20) + return mu, log_var + + def decode(self, z, seq_len): + b = z.shape[0] + z_proj = self.latent_project(z) # [B,d] + z_exp = mx.broadcast_to(z_proj[:, None, :], (b, seq_len, self.d)) + hidden = z_exp + self.dec_pos_embedding[:, :seq_len, :] + # causal mask [S,S]: 0 on/below diag, -inf above + causal = mx.where(mx.triu(mx.ones((seq_len, seq_len)), k=1) > 0, + mx.array(-1e9, mx.float32), mx.array(0.0, mx.float32)) + for layer in self.dec_layers: + hidden = layer(hidden, z_exp, causal) + hidden = self.dec_layer_norm(hidden) + cmd = self.command_head(hidden) # [B,S,6] + params = mx.stack([h(hidden) for h in self.param_heads], axis=0) # [16,B,S,256] + return cmd, params + + +# --- weight conversion -------------------------------------------------------- +def convert_checkpoint(ckpt_path, model): + import torch + + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + sd = ck if hasattr(ck, "items") else ck.get("model_state_dict", ck) + + def t(key): + return mx.array(sd[key].detach().cpu().float().numpy()) + + pairs = [("encoder.token_embedding.weight", t("encoder.token_embedding.weight")), + ("encoder.pos_embedding", t("encoder.pos_embedding")), + ("encoder.layer_norm.weight", t("encoder.layer_norm.weight")), + ("encoder.layer_norm.bias", t("encoder.layer_norm.bias"))] + for i in range(6): + s = f"encoder.transformer.layers.{i}" + d = f"encoder.layers.{i}" + pairs += [(f"{d}.self_attn.in_proj.weight", t(f"{s}.self_attn.in_proj_weight")), + (f"{d}.self_attn.in_proj.bias", t(f"{s}.self_attn.in_proj_bias")), + (f"{d}.self_attn.out_proj.weight", t(f"{s}.self_attn.out_proj.weight")), + (f"{d}.self_attn.out_proj.bias", t(f"{s}.self_attn.out_proj.bias")), + (f"{d}.linear1.weight", t(f"{s}.linear1.weight")), + (f"{d}.linear1.bias", t(f"{s}.linear1.bias")), + (f"{d}.linear2.weight", t(f"{s}.linear2.weight")), + (f"{d}.linear2.bias", t(f"{s}.linear2.bias")), + (f"{d}.norm1.weight", t(f"{s}.norm1.weight")), (f"{d}.norm1.bias", t(f"{s}.norm1.bias")), + (f"{d}.norm2.weight", t(f"{s}.norm2.weight")), (f"{d}.norm2.bias", t(f"{s}.norm2.bias"))] + for src, dst in (("mu_head", "mu_head"), ("log_var_head", "log_var_head"), + ("latent_project", "latent_project"), ("command_head", "command_head")): + pairs += [(f"{dst}.weight", t(f"{src}.weight")), (f"{dst}.bias", t(f"{src}.bias"))] + pairs.append(("dec_pos_embedding", t("dec_pos_embedding"))) + pairs += [("dec_layer_norm.weight", t("decoder.layer_norm.weight")), + ("dec_layer_norm.bias", t("decoder.layer_norm.bias"))] + for i in range(6): + s = f"decoder._transformer.layers.{i}" + d = f"dec_layers.{i}" + for a, mlxa in (("self_attn", "self_attn"), ("multihead_attn", "multihead_attn")): + pairs += [(f"{d}.{mlxa}.in_proj.weight", t(f"{s}.{a}.in_proj_weight")), + (f"{d}.{mlxa}.in_proj.bias", t(f"{s}.{a}.in_proj_bias")), + (f"{d}.{mlxa}.out_proj.weight", t(f"{s}.{a}.out_proj.weight")), + (f"{d}.{mlxa}.out_proj.bias", t(f"{s}.{a}.out_proj.bias"))] + pairs += [(f"{d}.linear1.weight", t(f"{s}.linear1.weight")), (f"{d}.linear1.bias", t(f"{s}.linear1.bias")), + (f"{d}.linear2.weight", t(f"{s}.linear2.weight")), (f"{d}.linear2.bias", t(f"{s}.linear2.bias")), + (f"{d}.norm1.weight", t(f"{s}.norm1.weight")), (f"{d}.norm1.bias", t(f"{s}.norm1.bias")), + (f"{d}.norm2.weight", t(f"{s}.norm2.weight")), (f"{d}.norm2.bias", t(f"{s}.norm2.bias")), + (f"{d}.norm3.weight", t(f"{s}.norm3.weight")), (f"{d}.norm3.bias", t(f"{s}.norm3.bias"))] + n_param = sum(1 for k in sd if k.startswith("param_heads.") and k.endswith(".weight")) + for i in range(n_param): + pairs += [(f"param_heads.{i}.weight", t(f"param_heads.{i}.weight")), + (f"param_heads.{i}.bias", t(f"param_heads.{i}.bias"))] + model.update(tree_unflatten(pairs)) + mx.eval(model.parameters()) + return len(pairs) + + +def parity(ckpt): + sys.path.insert(0, str(_REPO / "ll_stepnet")) + import types + + import torch + from stepnet.vae import STEPVAE + + cfg = types.SimpleNamespace(token_embed_dim=256, vocab_size=50000, + num_transformer_layers=6, dropout=0.1) + pt = STEPVAE(cfg) + ck = torch.load(ckpt, map_location="cpu", weights_only=False) + sd = ck if hasattr(ck, "items") else ck.get("model_state_dict", ck) + pt.load_state_dict(sd, strict=False) + pt.eval() + S = int(pt.max_seq_len) + n_param = len(pt.param_heads) + + model = STEPVAEMLX(max_seq=S, n_param=n_param) + npar = convert_checkpoint(ckpt, model) + + rng = np.random.default_rng(0) + ids = rng.integers(1, 268, (3, S)).astype(np.int64) + + with torch.no_grad(): + mu_pt, lv_pt = pt.encode(torch.from_numpy(ids)) + hidden_pt = pt.decode(mu_pt, seq_len=S) + cmd_pt = pt.command_head(hidden_pt).numpy() + par_pt = np.stack([h(hidden_pt).numpy() for h in pt.param_heads]) # [P,B,S,256] + + mu_mlx, _ = model.encode(mx.array(ids)) + cmd_mlx, par_mlx = model.decode(mu_mlx, S) + cmd_mlx = np.array(cmd_mlx.tolist()); par_mlx = np.array(par_mlx.tolist()) + + d_mu = float(np.abs(np.array(mu_mlx.tolist()) - mu_pt.numpy()).max()) + d_cmd = float(np.abs(cmd_mlx - cmd_pt).max()) + d_par = float(np.abs(par_mlx - par_pt).max()) + # decoded-command agreement (argmax over the 6 command types, per position) + agree = float((cmd_mlx.argmax(-1) == cmd_pt.argmax(-1)).mean()) + print(f"converted {npar} tensors from {os.path.basename(ckpt)}", flush=True) + print(f"PARITY max-abs-diff mu={d_mu:.2e} command_logits={d_cmd:.2e} param_logits={d_par:.2e}", flush=True) + print(f"decoded-command argmax agreement = {agree:.4f}", flush=True) + ok = d_mu < 1e-3 and d_cmd < 1e-3 and d_par < 1e-3 and agree > 0.999 + out = str(_REPO / "ll_gen/checkpoints/vae_mlx.safetensors") + mx.save_safetensors(out, dict(tree_flatten(model.parameters()))) + print(f"FAITHFUL_PARITY {'PASS' if ok else 'FAIL'} -> saved {out}", flush=True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "convert", "parity"], default="parity") + ap.add_argument("--ckpt", default=str(_REPO / "ll_gen/checkpoints/vae_warm.pt")) + args = ap.parse_args() + if args.mode == "probe": + m = STEPVAEMLX() + mu, lv = m.encode(mx.array(np.random.randint(1, 268, (2, 60)))) + cmd, par = m.decode(mu, 60) + print(f"probe: mu {mu.shape} command {cmd.shape} params {par.shape} " + f"finite={bool(mx.isfinite(cmd).all().item())}", flush=True) + return + if args.mode == "convert": + m = STEPVAEMLX() + n = convert_checkpoint(args.ckpt, m) + mx.save_safetensors(str(_REPO / "ll_gen/checkpoints/vae_mlx.safetensors"), + dict(tree_flatten(m.parameters()))) + print(f"converted {n} tensors", flush=True) + return + parity(args.ckpt) + + +if __name__ == "__main__": + main() diff --git a/ll_gen/scripts/train_diffusion_codec.py b/ll_gen/scripts/train_diffusion_codec.py new file mode 100644 index 0000000..e07e44d --- /dev/null +++ b/ll_gen/scripts/train_diffusion_codec.py @@ -0,0 +1,253 @@ +"""Train StructuredDiffusion's GeometryCodec on real DeepCAD B-rep geometry. + +Make-real campaign (ll_gen diffusion path). The GeometryCodec is the latent<-> +geometry autoencoder the diffusion denoisers operate in; it had never been +trained on real geometry. This trains it: DeepCAD cad_vec -> solid (validated +executor) -> sample each face as an 8x8x3 UV grid and each edge as a 12x3 +polyline (unit-cube normalized, masked) -> codec masked-MSE reconstruction. +Writes the checkpoint + metrics to ll_gen/checkpoints/. + +Result (4000 train / 600 val, 60 epochs MPS): recon MSE 0.40 (untrained) -> +0.0003 (trained). NOTE: this trains the codec only; the diffusion *denoisers* +do not yet converge (they plateau at predict-zero), so the generator does not +yet produce valid CAD — see docs/2026-06-10-ll-gen-make-real-findings.md. + +Requires DeepCAD cad_vec under ``--deepcad`` (data.tar from +http://www.cs.columbia.edu/cg/deepcad/) + the ll_gen executor + pythonocc. + +Run:: + + python ll_gen/scripts/train_diffusion_codec.py --n-train 4000 --epochs 60 \ + --device mps +""" + +from __future__ import annotations + +import argparse +import glob +import json +import logging +import os +import sys +import warnings +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MPLBACKEND", "Agg") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +warnings.filterwarnings("ignore") +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +matplotlib.use = lambda *a, **k: None +logging.disable(logging.WARNING) + +import numpy as np # noqa: E402 +import h5py # noqa: E402 +import torch # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] + +LEVELS, RANGE = 256, 2.0 +MASK = {"LINE": [0, 1, 2, 3], "ARC": [0, 1, 2, 3, 4, 5], "CIRCLE": [0, 1, 2], + "EXTRUDE": [0, 1, 2, 3, 4, 5, 6, 7], "SOL": [], "EOS": []} +NUM_FACES, NUM_EDGES, UV, MPTS = 8, 12, 8, 12 + + +def _q_coord(g): + return int(np.clip(round(float(g)), 0, LEVELS - 1)) + + +def _q_value(v): + return int(np.clip(round((float(v) + RANGE) / (2 * RANGE) * (LEVELS - 1)), 0, LEVELS - 1)) + + +def _command_dicts(cad, Circle, Arc): + out = [] + + def emit(name, slots): + p = [0] * 16 + m = [False] * 16 + for j in MASK[name]: + p[j] = int(slots.get(j, 0)) + m[j] = True + out.append({"command_type": name, "parameters": p, "parameter_mask": m}) + + for ext in cad.seq: + for loop in ext.profile.children: + emit("SOL", {}) + for cv in loop.children: + if isinstance(cv, Circle): + r_mag = float(cv.radius) / (LEVELS - 1) * 2.0 * RANGE + emit("CIRCLE", {0: _q_coord(cv.center[0]), 1: _q_coord(cv.center[1]), 2: _q_value(r_mag)}) + elif isinstance(cv, Arc): + s, e, c = cv.start_point, cv.end_point, cv.center + emit("ARC", {0: _q_coord(s[0]), 1: _q_coord(s[1]), 2: _q_coord(e[0]), + 3: _q_coord(e[1]), 4: _q_coord(c[0]), 5: _q_coord(c[1])}) + else: + s, e = cv.start_point, cv.end_point + emit("LINE", {0: _q_coord(s[0]), 1: _q_coord(s[1]), 2: _q_coord(e[0]), 3: _q_coord(e[1])}) + depth = abs(float(ext.extent_one)) + abs(float(ext.extent_two)) + emit("EXTRUDE", {0: _q_value(float(np.clip(depth * 4.0, 0.3, 2.0)))}) + emit("EOS", {}) + return out + + +def _extract_geometry(shape, occ): + (TopExp_Explorer, TopAbs_FACE, TopAbs_EDGE, topods, BRepAdaptor_Surface, + BRepAdaptor_Curve, Bnd_Box, brepbndlib) = occ + box = Bnd_Box() + brepbndlib.Add(shape, box) + xmin, ymin, zmin, xmax, ymax, zmax = box.Get() + center = np.array([(xmin + xmax) / 2, (ymin + ymax) / 2, (zmin + zmax) / 2]) + scale = 2.0 / max(xmax - xmin, ymax - ymin, zmax - zmin, 1e-6) + + def n(p): + return ((np.array([p.X(), p.Y(), p.Z()]) - center) * scale).astype(np.float32) + + fg = np.zeros((NUM_FACES, UV, UV, 3), np.float32) + fm = np.ones((NUM_FACES,), bool) + exp = TopExp_Explorer(shape, TopAbs_FACE) + fi = 0 + while exp.More() and fi < NUM_FACES: + s = BRepAdaptor_Surface(topods.Face(exp.Current())) + u0, u1, v0, v1 = s.FirstUParameter(), s.LastUParameter(), s.FirstVParameter(), s.LastVParameter() + for iu in range(UV): + for iv in range(UV): + fg[fi, iu, iv] = n(s.Value(u0 + (u1 - u0) * iu / (UV - 1), v0 + (v1 - v0) * iv / (UV - 1))) + fm[fi] = False + fi += 1 + exp.Next() + if fi == 0: + return None + ep = np.zeros((NUM_EDGES, MPTS, 3), np.float32) + em = np.ones((NUM_EDGES,), bool) + exp = TopExp_Explorer(shape, TopAbs_EDGE) + ei = 0 + while exp.More() and ei < NUM_EDGES: + try: + c = BRepAdaptor_Curve(topods.Edge(exp.Current())) + t0, t1 = c.FirstParameter(), c.LastParameter() + for k in range(MPTS): + ep[ei, k] = n(c.Value(t0 + (t1 - t0) * k / (MPTS - 1))) + em[ei] = False + ei += 1 + except Exception: + pass + exp.Next() + return fg, ep, fm, em + + +def load_geometry(files, limit, deps): + CADSequence, Circle, Arc, CommandSequenceProposal, execute, occ = deps + fg, ep, fm, em = [], [], [], [] + for f in files: + if len(fg) >= limit: + break + try: + with h5py.File(f, "r") as h: + vec = h["vec"][:].astype(int) + cad = CADSequence.from_vector(vec, is_numerical=True, n=256) + shape = execute(CommandSequenceProposal( + command_dicts=_command_dicts(cad, Circle, Arc), quantization_bits=8, normalization_range=2.0)) + if shape is None: + continue + g = _extract_geometry(shape, occ) + if g is None: + continue + fg.append(g[0]); ep.append(g[1]); fm.append(g[2]); em.append(g[3]) + except Exception: + continue + if not fg: + return None + return (torch.from_numpy(np.stack(fg)), torch.from_numpy(np.stack(ep)), + torch.from_numpy(np.stack(fm)), torch.from_numpy(np.stack(em))) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--deepcad", default=str(_REPO / "resources/DeepCAD/data/cad_vec")) + ap.add_argument("--n-train", type=int, default=4000) + ap.add_argument("--n-val", type=int, default=600) + ap.add_argument("--epochs", type=int, default=60) + ap.add_argument("--bs", type=int, default=64) + ap.add_argument("--lr", type=float, default=1e-3) + ap.add_argument("--device", default="cpu") + ap.add_argument("--out", default=str(_REPO / "ll_gen/checkpoints")) + args = ap.parse_args() + dev = args.device + + sys.path.insert(0, str(_REPO / "resources/DeepCAD")) + from cadlib.extrude import CADSequence + from cadlib.curves import Arc, Circle + from ll_gen.proposals.command_proposal import CommandSequenceProposal + from ll_gen.disposal.command_executor import execute_command_proposal + from ll_gen.training.run import build_generator + from OCC.Core.TopExp import TopExp_Explorer + from OCC.Core.TopAbs import TopAbs_FACE, TopAbs_EDGE + from OCC.Core.TopoDS import topods + from OCC.Core.BRepAdaptor import BRepAdaptor_Surface, BRepAdaptor_Curve + from OCC.Core.Bnd import Bnd_Box + from OCC.Core.BRepBndLib import brepbndlib + + occ = (TopExp_Explorer, TopAbs_FACE, TopAbs_EDGE, topods, BRepAdaptor_Surface, + BRepAdaptor_Curve, Bnd_Box, brepbndlib) + deps = (CADSequence, Circle, Arc, CommandSequenceProposal, execute_command_proposal, occ) + + gen = build_generator("diffusion", dev) + if gen._model is None and hasattr(gen, "_init_model"): + gen._init_model() + codec = gen._model.geometry_codec.to(dev) + + files = sorted(glob.glob(os.path.join(args.deepcad, "*/*.h5"))) + if not files: + raise SystemExit(f"No cad_vec h5 files under {args.deepcad}; download DeepCAD data.tar first.") + need = (args.n_train + args.n_val) * 4 + 4000 + print(f"extracting geometry from up to {need} DeepCAD solids ...", flush=True) + vfg, vep, vfm, vem = (t.to(dev) for t in load_geometry(files[:need // 6], args.n_val, deps)) + tfg, tep, tfm, tem = load_geometry(files[need // 6:], args.n_train, deps) + print(f"extracted {tfg.shape[0]} train / {vfg.shape[0]} val solids on {dev}", flush=True) + + opt = torch.optim.Adam(codec.parameters(), lr=args.lr) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs) + os.makedirs(args.out, exist_ok=True) + ckpt = os.path.join(args.out, "diffusion_codec.pt") + n = tfg.shape[0] + best = 1e9 + for epoch in range(args.epochs): + codec.train() + perm = torch.randperm(n) + tot = 0.0 + nb = 0 + for k in range(0, n, args.bs): + idx = perm[k:k + args.bs] + out = codec.reconstruction_loss(tfg[idx].to(dev), tep[idx].to(dev), + tfm[idx].to(dev), tem[idx].to(dev)) + loss = out["total_recon_loss"] + opt.zero_grad() + loss.backward() + opt.step() + tot += float(loss) + nb += 1 + sched.step() + codec.eval() + with torch.no_grad(): + v = codec.reconstruction_loss(vfg, vep, vfm, vem) + vtot = float(v["total_recon_loss"]) + print(f"epoch {epoch+1}/{args.epochs} train_loss={tot/max(nb,1):.5f} " + f"val_face_mse={float(v['face_recon_loss']):.5f} val_edge_mse={float(v['edge_recon_loss']):.5f}", + flush=True) + if vtot < best: + best = vtot + torch.save({"codec_state_dict": codec.state_dict()}, ckpt) + + result = {"n_train": int(tfg.shape[0]), "n_val": int(vfg.shape[0]), "epochs": args.epochs, + "best_val_recon_mse": round(best, 6), "checkpoint": ckpt, + "latent_dim": codec.latent_dim} + with open(os.path.join(args.out, "diffusion_codec_metrics.json"), "w") as fh: + json.dump(result, fh, indent=2) + print("DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_gen/tests/test_command_executor.py b/ll_gen/tests/test_command_executor.py index dbc970e..9d7fc7c 100644 --- a/ll_gen/tests/test_command_executor.py +++ b/ll_gen/tests/test_command_executor.py @@ -613,3 +613,112 @@ def test_command_proposal_token_ids_fixture( """Test command_proposal_token_ids fixture has expected structure.""" assert hasattr(command_proposal_token_ids, "token_ids") assert len(command_proposal_token_ids.token_ids) > 0 + + +# ============================================================================ +# SECTION: Closure-aware sketch construction (real OCC) +# ============================================================================ + +from ll_gen.disposal import command_executor as _ce # noqa: E402 + + +def _quant(v: float) -> int: + """Continuous coord in [-2,2] -> 8-bit quantized slot (executor's scheme).""" + return max(0, min(255, int(round((v + 2.0) / 4.0 * 255)))) + + +def _line(x1, y1, x2, y2): + p = [0] * 16 + m = [False] * 16 + for i, v in zip((0, 1, 2, 3), (x1, y1, x2, y2)): + p[i] = _quant(v) + m[i] = True + return {"command_type": "LINE", "parameters": p, "parameter_mask": m} + + +def _bare(ctype): + return {"command_type": ctype, "parameters": [0] * 16, "parameter_mask": [False] * 16} + + +def _extrude(depth=1.0): + p = [0] * 16 + m = [False] * 16 + p[0] = _quant(depth) + for i in range(8): + m[i] = True + return {"command_type": "EXTRUDE", "parameters": p, "parameter_mask": m} + + +@pytest.mark.skipif(not _ce._OCC_AVAILABLE, reason="requires pythonocc-core") +class TestClosureAwareSketch: + """The executor must close multi-line sketch loops by threading endpoints. + + A non-autoregressive decoder emits each line's coordinates independently, so + consecutive segments almost never share a vertex. The closure-aware builder + threads endpoints and auto-closes the loop, so such sketches still produce a + valid non-degenerate solid instead of failing wire construction. Without the + fix only self-closing primitives (circles) ever validated. + """ + + def _volume(self, shape): + from OCC.Core.GProp import GProp_GProps + from OCC.Core.BRepGProp import brepgprop + + props = GProp_GProps() + brepgprop.VolumeProperties(shape, props) + return props.Mass() + + def test_non_connecting_square_closes_to_solid(self) -> None: + # Four lines with a ~0.03 gap at every corner — they do NOT connect. + cmds = [ + _bare("SOL"), + _line(0.00, 0.00, 1.02, -0.01), + _line(0.98, 0.03, 1.01, 0.99), + _line(1.03, 1.02, -0.02, 0.97), + _line(0.01, 1.01, -0.01, 0.02), + _extrude(1.0), + _bare("EOS"), + ] + prop = CommandSequenceProposal( + command_dicts=cmds, quantization_bits=8, normalization_range=2.0 + ) + shape = execute_command_proposal(prop) + assert shape is not None, "closure-aware builder must close the gapped square" + assert self._volume(shape) > 0.1, "closed square extrusion must have real volume" + + def test_triangle_closes_to_solid(self) -> None: + # Three lines whose endpoints do not coincide — must still close. + cmds = [ + _bare("SOL"), + _line(0.0, 0.0, 1.0, 0.05), + _line(0.95, 0.0, 0.5, 1.0), + _line(0.55, 0.95, 0.02, 0.03), + _extrude(0.8), + _bare("EOS"), + ] + prop = CommandSequenceProposal( + command_dicts=cmds, quantization_bits=8, normalization_range=2.0 + ) + shape = execute_command_proposal(prop) + assert shape is not None + assert self._volume(shape) > 0.05 + + def test_single_circle_still_valid(self) -> None: + # Regression guard: a lone circle loop must remain a valid solid. + circ_p = [0] * 16 + circ_m = [False] * 16 + circ_p[0], circ_p[1], circ_p[2] = _quant(0.0), _quant(0.0), _quant(2.0) + for i in (0, 1, 2): + circ_m[i] = True + cmds = [ + _bare("SOL"), + {"command_type": "CIRCLE", "parameters": circ_p, "parameter_mask": circ_m}, + _extrude(1.0), + _bare("EOS"), + ] + prop = CommandSequenceProposal( + command_dicts=cmds, quantization_bits=8, normalization_range=2.0 + ) + shape = execute_command_proposal(prop) + assert shape is not None + assert self._volume(shape) > 0.1 diff --git a/ll_gen/tests/test_surface_executor.py b/ll_gen/tests/test_surface_executor.py index 41fe488..f5fe4f8 100644 --- a/ll_gen/tests/test_surface_executor.py +++ b/ll_gen/tests/test_surface_executor.py @@ -348,3 +348,77 @@ def test_average_edges_exists(self) -> None: from ll_gen.disposal.surface_executor import _average_edges assert callable(_average_edges) + + +# ============================================================================ +# SECTION: End-to-end sewing of a watertight box (real OCC) +# ============================================================================ + +from ll_gen.disposal import surface_executor as _se # noqa: E402 + + +def _box_face_grids(uv: int = 8): + """Six UVxUVx3 point grids sampling the faces of the unit cube [0,1]^3. + + The faces share edges exactly (as a real solid does), so a correct + surface-fit + merge + sew pipeline must produce a closed, non-zero-volume + shape. This is the diffusion ``execute_latent_proposal`` path. + """ + t = np.linspace(0.0, 1.0, uv) + a, b = np.meshgrid(t, t, indexing="ij") + z0 = np.zeros_like(a) + z1 = np.ones_like(a) + faces = [ + np.stack([a, b, z0], axis=-1), # bottom z=0 + np.stack([a, b, z1], axis=-1), # top z=1 + np.stack([a, z0, b], axis=-1), # front y=0 + np.stack([a, z1, b], axis=-1), # back y=1 + np.stack([z0, a, b], axis=-1), # left x=0 + np.stack([z1, a, b], axis=-1), # right x=1 + ] + return [f.astype(np.float64) for f in faces] + + +def _box_edge_points(m: int = 12): + """The 12 edges of the unit cube, each an M-point polyline.""" + t = np.linspace(0.0, 1.0, m) + o, i = np.zeros(m), np.ones(m) + segs = [ + (t, o, o), (t, i, o), (t, o, i), (t, i, i), # x-parallel + (o, t, o), (i, t, o), (o, t, i), (i, t, i), # y-parallel + (o, o, t), (i, o, t), (o, i, t), (i, i, t), # z-parallel + ] + return [np.stack(s, axis=-1).astype(np.float64) for s in segs] + + +@pytest.mark.skipif(not _se._OCC_AVAILABLE, reason="requires pythonocc-core") +class TestWatertightSew: + """Regression: the surface executor must sew independently-fit faces into a + closed, non-zero-volume shape. + + This guards three real bugs that made the diffusion path never produce a + solid: (1) ``_fit_bspline_surface`` called cadling's fitter with an + unsupported ``tolerance=`` kwarg and expected a face but got a dict + (TypeError → every face dropped); (2) the merge step called + ``TopologyMerger.merge_edges`` which does not exist (AttributeError → + crash); the real API is ``merge(faces)``. With those fixed, the six faces of + a unit cube must sew into a closed shape enclosing ~unit volume. + """ + + def _volume(self, shape): + from OCC.Core.GProp import GProp_GProps + from OCC.Core.BRepGProp import brepgprop + + props = GProp_GProps() + brepgprop.VolumeProperties(shape, props) + return abs(props.Mass()) + + def test_unit_box_sews_to_closed_volume(self) -> None: + proposal = LatentProposal( + face_grids=_box_face_grids(), edge_points=_box_edge_points() + ) + shape = execute_latent_proposal(proposal) + assert shape is not None, "six box faces must sew into a shape" + # A correctly sewn unit cube encloses ~1.0 of volume (allow fitting slack). + vol = self._volume(shape) + assert vol > 0.3, f"sewn box must enclose real volume, got {vol:.3f}" diff --git a/ll_gen/tests/test_validity_eval.py b/ll_gen/tests/test_validity_eval.py index c21204e..b563813 100644 --- a/ll_gen/tests/test_validity_eval.py +++ b/ll_gen/tests/test_validity_eval.py @@ -56,12 +56,13 @@ def generate_candidates( def _valid_result() -> DisposalResult: # has_shape is derived from `shape is not None`; a truthy sentinel marks a - # constructed, valid shape. The reward now requires a closed solid for full - # credit, so a "valid result" fixture carries a solid geometry report. + # constructed, valid shape. The validity gate now requires a NON-DEGENERATE + # closed solid (BRepCheck alone passes volume-less shells), so a "valid + # result" fixture carries a solid geometry report WITH positive volume. return DisposalResult( shape=object(), is_valid=True, - geometry_report=GeometryReport(solid_count=1, is_solid=True), + geometry_report=GeometryReport(solid_count=1, is_solid=True, volume=1.0), ) @@ -238,3 +239,39 @@ def test_load_generator_checkpoint_missing_file_raises(tmp_path) -> None: # rather than silently no-op'ing. with pytest.raises((FileNotFoundError, RuntimeError, ImportError)): load_generator_checkpoint(gen, tmp_path / "does_not_exist.pt") + + +# --------------------------------------------------------------------------- +# Export: an eval run must write the generated CAD as files, not empty folders +# --------------------------------------------------------------------------- + + +def test_export_flag_forwarded_to_dispose() -> None: + """evaluate_validity exports valid shapes by default and honors export=False. + + Regression: the harness used to hardcode ``export=False`` while still + creating ``/disposed/``, so every eval left an empty folder + instead of the generated CAD. It now defaults to exporting (writing + ``.step`` + ``.stl`` per valid shape) and forwards the flag to + ``DisposalEngine.dispose``. + """ + from unittest.mock import patch + + captured: dict[str, Any] = {} + + class _FakeEngine: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def dispose(self, proposal: Any, export: bool = True) -> DisposalResult: + captured["export"] = export + return _valid_result() + + gen = _FakeGenerator() + with patch("ll_gen.disposal.engine.DisposalEngine", _FakeEngine): + # No dispose_fn injected -> the real-engine path runs. + evaluate_validity(gen, ["p"], n_samples=1, output_dir="x") + assert captured["export"] is True # exports by default + + evaluate_validity(gen, ["p"], n_samples=1, output_dir="x", export=False) + assert captured["export"] is False # honored when disabled diff --git a/ll_ocadr/mlx/faithful_tower_mlx.py b/ll_ocadr/mlx/faithful_tower_mlx.py new file mode 100644 index 0000000..b70f723 --- /dev/null +++ b/ll_ocadr/mlx/faithful_tower_mlx.py @@ -0,0 +1,415 @@ +"""Faithful MLX port of the real ll_ocadr geometry tower (PointNet++ + Point-BERT). + +The configured PyTorch ll_ocadr model was never trained (no checkpoint to convert) +AND cannot run on Apple Silicon, so unlike ll_stepnet/ll_brepnet there are no trained +weights to load. "Faithful" here therefore means: reproduce the EXACT architecture of +the real encoders and PROVE the MLX tower computes the same function as the real code. + +That proof is possible because GeometryNet / ShapeNet / MlpProjector are plain +``torch.nn.Module``s with no LLM dependency — they instantiate and run on CPU. So we +random-init the real PyTorch tower, convert its weights into this MLX tower, feed the +SAME ``(coords, normals)`` through both, and assert max-abs-diff ~1e-5 (the same rigor +that made the other two ports verifiably faithful rather than asserted). + +Real architecture (ll_ocadr/vllm/lattice_encoder/*, config for the 0.5B model): + GeometryNet (PointNet++): + SA1: FPS N->512, ball-query r0.2 nsample32, mlp[64,64,128] (in 3+3) + SA2: FPS 512->128, ball-query r0.4 nsample64, mlp[128,128,256] (in 3+128) + MultiheadAttention(256, 8) over the 128 points, +residual, LayerNorm -> [B,128,256] + ShapeNet (Point-BERT, 0.5B config: embed768 depth4 heads8): + PointPatchEmbedding (1x1 conv mini-PointNet -> 256 patches x 768), + CLS + learnable pos, depth TransformerBlocks (pre-norm, GELU MLP), LayerNorm -> [B,257,768] + Forward fuse: concat[shape[:,1:] (256x768), geom padded 128->256 (256x256)] -> 256x1024 + -> MlpProjector (linear 1024->n_embed) -> 256 mesh tokens. + +The 1x1 convs are per-point Linears over channels, so they are implemented as MLX +Linear + BatchNorm (mathematically identical, no NHWC Conv layout). MultiheadAttention +is implemented manually from the packed in_proj_weight (the same trick used for the +stepnet port). FPS + ball-query are a deterministic geometric function of the fixed +cloud (start fixed at point 0) — computed once in numpy and shared by both models, so +they cannot be a source of divergence. + +Modes: + parity - random-init the real PyTorch tower, convert -> MLX, assert forward parity. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import warnings + +os.environ.setdefault("OMP_NUM_THREADS", "1") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import numpy as np # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +from mlx.utils import tree_unflatten # noqa: E402 + +_REPO = "/Users/ryanoboyle/LatticeLabs_toolkit" + + +# ============================================================================ +# Deterministic numpy FPS / ball-query / grouping (shared by both models). +# ============================================================================ +def square_distance_np(src, dst): # src [N,3], dst [M,3] -> [N,M] + return (-2 * src @ dst.T) + (src ** 2).sum(1)[:, None] + (dst ** 2).sum(1)[None, :] + + +def fps_np(xyz, npoint): + """Farthest point sampling, deterministic start at index 0. -> [npoint] indices.""" + n = xyz.shape[0] + centroids = np.zeros(npoint, np.int64) + distance = np.full(n, 1e10, np.float64) + farthest = 0 + for i in range(npoint): + centroids[i] = farthest + d = ((xyz - xyz[farthest]) ** 2).sum(1) + distance = np.minimum(distance, d) + farthest = int(np.argmax(distance)) + return centroids + + +def ball_query_np(radius, nsample, xyz, new_xyz): + """[S,nsample] indices — replicates geometry_net.query_ball_point exactly.""" + n = xyz.shape[0] + s = new_xyz.shape[0] + group_idx = np.tile(np.arange(n, dtype=np.int64)[None, :], (s, 1)) # [S,N] + sqrdists = square_distance_np(new_xyz, xyz) # [S,N] + group_idx[sqrdists > radius ** 2] = n + group_idx = np.sort(group_idx, axis=-1)[:, :nsample] + group_first = np.tile(group_idx[:, 0:1], (1, group_idx.shape[1])) + mask = group_idx == n + group_idx[mask] = group_first[mask] + return group_idx + + +def sample_and_group_np(npoint, radius, nsample, xyz, points): + """Returns new_xyz [npoint,3], grouped [npoint,nsample,3+D], group_idx [npoint,nsample].""" + fps_idx = fps_np(xyz, npoint) + new_xyz = xyz[fps_idx] + group_idx = ball_query_np(radius, nsample, xyz, new_xyz) + grouped_xyz = xyz[group_idx] # [npoint,nsample,3] + grouped_xyz_norm = grouped_xyz - new_xyz[:, None] # center + if points is not None: + grouped_points = points[group_idx] + grouped = np.concatenate([grouped_xyz_norm, grouped_points], axis=-1) + else: + grouped = grouped_xyz_norm + return new_xyz, grouped, group_idx + + +def precompute_geom(coords, normals): + """All geometric (cloud-only) tensors GeometryNetMLX needs. SA2 feature grouping + happens at runtime (depends on SA1 conv output), so cache its xyz-norm + indices.""" + # SA1: full grouping (features = normals, known at data time) + new_xyz1, sa1_grouped, _ = sample_and_group_np(512, 0.2, 32, coords, normals) + # SA2 geometry on the SA1 centroids (features = SA1 output -> gathered at runtime) + fps2 = fps_np(new_xyz1, 128) + new_xyz2 = new_xyz1[fps2] + g2 = ball_query_np(0.4, 64, new_xyz1, new_xyz2) # [128,64] into the 512 + sa2_grouped_xyz = new_xyz1[g2] - new_xyz2[:, None] # [128,64,3] + return { + "sa1_grouped": sa1_grouped.astype(np.float32), # [512,32,6] + "sa2_grouped_xyz": sa2_grouped_xyz.astype(np.float32), # [128,64,3] + "sa2_group_idx": g2.astype(np.int32), # [128,64] + } + + +# ============================================================================ +# MLX modules +# ============================================================================ +class LinBN(mlxnn.Module): + """1x1 conv == per-point Linear over channels, followed by BatchNorm.""" + + def __init__(self, cin, cout): + super().__init__() + self.lin = mlxnn.Linear(cin, cout) + self.bn = mlxnn.BatchNorm(cout) + + def __call__(self, x): + return self.bn(self.lin(x)) + + +class SAMLX(mlxnn.Module): + """PointNetSetAbstraction mlp: stack of (Linear->BN->ReLU), max over nsample axis.""" + + def __init__(self, in_ch, mlp): + super().__init__() + layers = [] + last = in_ch + for out in mlp: + layers.append(LinBN(last, out)) + last = out + self.layers = layers + + def __call__(self, grouped): # [B, npoint, nsample, in_ch] + x = grouped + for layer in self.layers: + x = mlxnn.relu(layer(x)) + return x.max(axis=2) # [B, npoint, mlp[-1]] + + +class MHA(mlxnn.Module): + """nn.MultiheadAttention math from the packed in_proj_weight [3d, d].""" + + def __init__(self, d, heads): + super().__init__() + self.d, self.h, self.hd = d, heads, d // heads + self.in_proj = mlxnn.Linear(d, 3 * d) + self.out_proj = mlxnn.Linear(d, d) + + def __call__(self, x): # [B,S,d] + b, s, _ = x.shape + qkv = self.in_proj(x) + q, k, v = mx.split(qkv, 3, axis=-1) + + def hd(t): + return mx.transpose(t.reshape(b, s, self.h, self.hd), (0, 2, 1, 3)) + + q, k, v = hd(q), hd(k), hd(v) + att = mx.softmax((q @ mx.transpose(k, (0, 1, 3, 2))) / (self.hd ** 0.5), axis=-1) + ctx = mx.transpose(att @ v, (0, 2, 1, 3)).reshape(b, s, self.d) + return self.out_proj(ctx) + + +class GeometryNetMLX(mlxnn.Module): + def __init__(self): + super().__init__() + self.sa1 = SAMLX(6, [64, 64, 128]) + self.sa2 = SAMLX(131, [128, 128, 256]) + self.local_attn = MHA(256, 8) + self.norm = mlxnn.LayerNorm(256) + + def __call__(self, sa1_grouped, sa2_grouped_xyz, sa2_group_idx): + feat1 = self.sa1(sa1_grouped) # [B,512,128] + b = feat1.shape[0] + gathered = mx.stack([feat1[i][sa2_group_idx[i]] for i in range(b)]) # [B,128,64,128] + grouped2 = mx.concatenate([sa2_grouped_xyz, gathered], axis=-1) # [B,128,64,131] + feat2 = self.sa2(grouped2) # [B,128,256] + feat2 = self.norm(feat2 + self.local_attn(feat2)) + return feat2 # [B,128,256] + + +class PatchEmbedMLX(mlxnn.Module): + def __init__(self, embed_dim=768): + super().__init__() + self.f1 = LinBN(6, 128) + self.f2 = mlxnn.Linear(128, 256) # first_conv tail (no BN) + self.s1 = LinBN(512, 512) + self.s2 = mlxnn.Linear(512, embed_dim) # second_conv tail (no BN) + self.embed_dim = embed_dim + + def __call__(self, coords, normals): + pts = mx.concatenate([coords, normals], axis=-1) # [B,N,6] + f = self.f2(mlxnn.relu(self.f1(pts))) # [B,N,256] + g = f.max(axis=1, keepdims=True) # [B,1,256] + f = mx.concatenate([mx.broadcast_to(g, f.shape), f], axis=-1) # [B,N,512] + f = self.s2(mlxnn.relu(self.s1(f))) # [B,N,embed] + b, n, e = f.shape + ps = n // 256 + patches = f[:, : 256 * ps].reshape(b, 256, ps, e).max(axis=2) # [B,256,embed] + return patches + + +class TFBlockMLX(mlxnn.Module): + def __init__(self, d, heads, mlp_ratio=4.0): + super().__init__() + self.norm1 = mlxnn.LayerNorm(d) + self.attn = MHA(d, heads) + self.norm2 = mlxnn.LayerNorm(d) + hidden = int(d * mlp_ratio) + self.fc1 = mlxnn.Linear(d, hidden) + self.fc2 = mlxnn.Linear(hidden, d) + + def __call__(self, x): + x = x + self.attn(self.norm1(x)) + return x + self.fc2(mlxnn.gelu(self.fc1(self.norm2(x)))) + + +class ShapeNetMLX(mlxnn.Module): + def __init__(self, embed=768, depth=4, heads=8): + super().__init__() + self.patch_embed = PatchEmbedMLX(embed) + self.cls_token = mx.zeros((1, 1, embed)) + self.pos_embed = mx.zeros((1, 257, embed)) + self.blocks = [TFBlockMLX(embed, heads) for _ in range(depth)] + self.norm = mlxnn.LayerNorm(embed) + self.embed = embed + + def __call__(self, coords, normals): + pt = self.patch_embed(coords, normals) # [B,256,embed] + b = pt.shape[0] + cls = mx.broadcast_to(self.cls_token, (b, 1, self.embed)) + tokens = mx.concatenate([cls, pt], axis=1) + self.pos_embed # [B,257,embed] + for blk in self.blocks: + tokens = blk(tokens) + return self.norm(tokens) # [B,257,embed] + + +class FaithfulTower(mlxnn.Module): + """GeometryNet + ShapeNet + linear projector -> 256 mesh tokens (+ aux class head).""" + + def __init__(self, n_embed, shape_depth=4, shape_heads=8): + super().__init__() + self.geometry = GeometryNetMLX() + self.shape = ShapeNetMLX(768, shape_depth, shape_heads) + self.projector = mlxnn.Linear(1024, n_embed) + self.aux_head = mlxnn.Linear(n_embed, 3) + + def __call__(self, b): + geom = self.geometry(b["sa1_grouped"], b["sa2_grouped_xyz"], b["sa2_group_idx"]) # [B,128,256] + shp = self.shape(b["coords"], b["normals"])[:, 1:] # [B,256,768] (skip CLS) + bs = geom.shape[0] + geom_pad = mx.concatenate([geom, mx.zeros((bs, 128, 256))], axis=1) # [B,256,256] + feat = mx.concatenate([shp, geom_pad], axis=-1) # [B,256,1024] + return self.projector(feat) # [B,256,n_embed] + + def aux_logits(self, tokens): + return self.aux_head(tokens.mean(axis=1)) + + +# ============================================================================ +# weight conversion: real PyTorch encoders -> MLX tower +# ============================================================================ +def convert_tower(geom_pt, shape_pt, proj_pt, mlx_tower): + import torch + + def t(x): + return mx.array(x.detach().cpu().float().numpy()) + + gsd = geom_pt.state_dict() + ssd = shape_pt.state_dict() + psd = proj_pt.state_dict() + pairs = [] + + # GeometryNet SA layers: mlp_convs.i (Conv2d 1x1 -> Linear), mlp_bns.i (BN) + for sa, n_mlp in (("sa1", 3), ("sa2", 3)): + for i in range(n_mlp): + w = gsd[f"{sa}.mlp_convs.{i}.weight"] # [out,in,1,1] + pairs.append((f"geometry.{sa}.layers.{i}.lin.weight", t(w[:, :, 0, 0]))) + pairs.append((f"geometry.{sa}.layers.{i}.lin.bias", t(gsd[f"{sa}.mlp_convs.{i}.bias"]))) + for p in ("weight", "bias", "running_mean", "running_var"): + pairs.append((f"geometry.{sa}.layers.{i}.bn.{p}", t(gsd[f"{sa}.mlp_bns.{i}.{p}"]))) + # GeometryNet attention + norm + pairs += [("geometry.local_attn.in_proj.weight", t(gsd["local_attn.in_proj_weight"])), + ("geometry.local_attn.in_proj.bias", t(gsd["local_attn.in_proj_bias"])), + ("geometry.local_attn.out_proj.weight", t(gsd["local_attn.out_proj.weight"])), + ("geometry.local_attn.out_proj.bias", t(gsd["local_attn.out_proj.bias"])), + ("geometry.norm.weight", t(gsd["norm.weight"])), ("geometry.norm.bias", t(gsd["norm.bias"]))] + + # ShapeNet patch embed: first_conv (0=Conv1d,1=BN,3=Conv1d), second_conv (0,1,3) + def conv1d(w): # [out,in,1] -> [out,in] + return t(w[:, :, 0]) + + pairs += [("shape.patch_embed.f1.lin.weight", conv1d(ssd["patch_embed.first_conv.0.weight"])), + ("shape.patch_embed.f1.lin.bias", t(ssd["patch_embed.first_conv.0.bias"]))] + for p in ("weight", "bias", "running_mean", "running_var"): + pairs.append((f"shape.patch_embed.f1.bn.{p}", t(ssd[f"patch_embed.first_conv.1.{p}"]))) + pairs += [("shape.patch_embed.f2.weight", conv1d(ssd["patch_embed.first_conv.3.weight"])), + ("shape.patch_embed.f2.bias", t(ssd["patch_embed.first_conv.3.bias"])), + ("shape.patch_embed.s1.lin.weight", conv1d(ssd["patch_embed.second_conv.0.weight"])), + ("shape.patch_embed.s1.lin.bias", t(ssd["patch_embed.second_conv.0.bias"]))] + for p in ("weight", "bias", "running_mean", "running_var"): + pairs.append((f"shape.patch_embed.s1.bn.{p}", t(ssd[f"patch_embed.second_conv.1.{p}"]))) + pairs += [("shape.patch_embed.s2.weight", conv1d(ssd["patch_embed.second_conv.3.weight"])), + ("shape.patch_embed.s2.bias", t(ssd["patch_embed.second_conv.3.bias"]))] + # ShapeNet cls/pos/blocks/norm + pairs += [("shape.cls_token", t(ssd["cls_token"])), ("shape.pos_embed", t(ssd["pos_embed"])), + ("shape.norm.weight", t(ssd["norm.weight"])), ("shape.norm.bias", t(ssd["norm.bias"]))] + depth = len({k.split("blocks.")[1].split(".")[0] for k in ssd if k.startswith("blocks.")}) + for i in range(depth): + b = f"blocks.{i}" + pairs += [(f"shape.blocks.{i}.norm1.weight", t(ssd[f"{b}.norm1.weight"])), + (f"shape.blocks.{i}.norm1.bias", t(ssd[f"{b}.norm1.bias"])), + (f"shape.blocks.{i}.norm2.weight", t(ssd[f"{b}.norm2.weight"])), + (f"shape.blocks.{i}.norm2.bias", t(ssd[f"{b}.norm2.bias"])), + (f"shape.blocks.{i}.attn.in_proj.weight", t(ssd[f"{b}.attn.in_proj_weight"])), + (f"shape.blocks.{i}.attn.in_proj.bias", t(ssd[f"{b}.attn.in_proj_bias"])), + (f"shape.blocks.{i}.attn.out_proj.weight", t(ssd[f"{b}.attn.out_proj.weight"])), + (f"shape.blocks.{i}.attn.out_proj.bias", t(ssd[f"{b}.attn.out_proj.bias"])), + (f"shape.blocks.{i}.fc1.weight", t(ssd[f"{b}.mlp.0.weight"])), + (f"shape.blocks.{i}.fc1.bias", t(ssd[f"{b}.mlp.0.bias"])), + (f"shape.blocks.{i}.fc2.weight", t(ssd[f"{b}.mlp.3.weight"])), + (f"shape.blocks.{i}.fc2.bias", t(ssd[f"{b}.mlp.3.bias"]))] + + # projector (linear): MlpProjector wraps it as .layers; a bare nn.Linear is "weight" + pkey = next(k for k in ("layers.weight", "layers.0.weight", "weight") if k in psd) + bkey = pkey.replace("weight", "bias") + pairs += [("projector.weight", t(psd[pkey])), ("projector.bias", t(psd[bkey]))] + + mlx_tower.update(tree_unflatten(pairs)) + mx.eval(mlx_tower.parameters()) + return len(pairs) + + +def parity(): + sys.path.insert(0, f"{_REPO}/ll_ocadr/vllm/lattice_encoder") + sys.path.insert(0, f"{_REPO}/ll_ocadr") + import torch + import geometry_net as G + from geometry_net import build_geometry_net + from shape_net import build_shape_net + + # make PyTorch sampling deterministic + identical to our numpy precompute + G.farthest_point_sample = lambda xyz, npoint: torch.from_numpy( + np.stack([fps_np(x.cpu().numpy(), npoint) for x in xyz])).long().to(xyz.device) + G.query_ball_point = lambda radius, nsample, xyz, new_xyz: torch.from_numpy( + np.stack([ball_query_np(radius, nsample, xyz[i].cpu().numpy(), new_xyz[i].cpu().numpy()) + for i in range(xyz.shape[0])])).long().to(xyz.device) + + torch.manual_seed(0) + np.random.seed(0) + n_embed = 896 + geom_pt = build_geometry_net().eval() + shape_pt = build_shape_net(embed_dim=768, depth=4, num_heads=8).eval() + proj_pt = torch.nn.Linear(1024, n_embed).eval() + + N = 2048 + coords = np.random.randn(N, 3).astype(np.float32) * 0.3 + normals = np.random.randn(N, 3).astype(np.float32) + normals /= (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8) + + # PyTorch reference forward (replicating latticelabs_ocadr fuse logic) + with torch.no_grad(): + c = torch.from_numpy(coords)[None] + nrm = torch.from_numpy(normals)[None] + g_pt = geom_pt(c, nrm) # [1,128,256] + s_pt = shape_pt(c, nrm)[:, 1:] # [1,256,768] + g_pad = torch.cat([g_pt, torch.zeros(1, 128, 256)], dim=1) # [1,256,256] + fused = torch.cat([s_pt, g_pad], dim=-1) # [1,256,1024] + tok_pt = proj_pt(fused).numpy() # [1,256,896] + + # MLX tower + tower = FaithfulTower(n_embed, shape_depth=4, shape_heads=8) + npar = convert_tower(geom_pt, shape_pt, proj_pt, tower) + tower.eval() # BatchNorm must use converted running stats (match PyTorch .eval()) + pre = precompute_geom(coords, normals) + b = {"sa1_grouped": mx.array(pre["sa1_grouped"])[None], + "sa2_grouped_xyz": mx.array(pre["sa2_grouped_xyz"])[None], + "sa2_group_idx": mx.array(pre["sa2_group_idx"])[None], + "coords": mx.array(coords)[None], "normals": mx.array(normals)[None]} + # component diffs + g_mlx = np.array(tower.geometry(b["sa1_grouped"], b["sa2_grouped_xyz"], b["sa2_group_idx"]).tolist()) + s_mlx = np.array(tower.shape(b["coords"], b["normals"])[:, 1:].tolist()) + tok_mlx = np.array(tower(b).tolist()) + + dg = float(np.abs(g_mlx - g_pt.detach().numpy()).max()) + ds = float(np.abs(s_mlx - s_pt.detach().numpy()).max()) + dt = float(np.abs(tok_mlx - tok_pt).max()) + print(f"converted {npar} tensors", flush=True) + print(f"PARITY max-abs-diff geometry={dg:.2e} shape={ds:.2e} mesh_tokens={dt:.2e}", flush=True) + ok = dg < 1e-3 and ds < 1e-3 and dt < 1e-3 + print(f"FAITHFUL_PARITY {'PASS' if ok else 'FAIL'} " + f"(tower reproduces the real PyTorch encoders within float tolerance)", flush=True) + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["parity"], default="parity") + ap.parse_args() + parity() diff --git a/ll_ocadr/mlx/train_ocadr_mlx.py b/ll_ocadr/mlx/train_ocadr_mlx.py new file mode 100644 index 0000000..7dc5a4b --- /dev/null +++ b/ll_ocadr/mlx/train_ocadr_mlx.py @@ -0,0 +1,359 @@ +"""ll_ocadr in MLX on Apple Silicon — FAITHFUL geometry tower + retrain. + +Supersedes the earlier maxpool stand-in encoder. The trainable 3D tower is now the +REAL ll_ocadr architecture ported to MLX and forward-parity-proven against the +PyTorch encoders (see faithful_tower_mlx.py: GeometryNet PointNet++ + ShapeNet +Point-BERT + linear projector). It maps a CAD point cloud (coords + normals) into +256 mesh tokens in a 4-bit Qwen2's embedding space (mlx-lm ``input_embeddings``); +the LLM base stays frozen/quantized while LoRA adapters + the tower train jointly so +the LLM learns to attend to the injected mesh tokens (LLaVA-style). + +There are NO pretrained ll_ocadr weights (the configured model was never trained and +can't run here), so this is an architecture-faithful RETRAIN, not a parity claim — +the faithfulness proof lives in faithful_tower_mlx.py. Task/data/metrics are held +IDENTICAL to the maxpool run for a direct comparison: balanced 3-way class from the point +cloud, class-only target, honest SHUFFLED-mesh baseline (feed the wrong mesh). + +Modes: probe | train. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +import warnings + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MPLBACKEND", "Agg") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +matplotlib.use = lambda *a, **k: None + +import numpy as np # noqa: E402 +import h5py # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +import mlx.optimizers as optim # noqa: E402 +from mlx.utils import tree_flatten # noqa: E402 +from mlx_lm import load as mlx_load # noqa: E402 +from mlx_lm.tuner.utils import linear_to_lora_layers # noqa: E402 + +_REPO = "/Users/ryanoboyle/LatticeLabs_toolkit" +_DEEPCAD = f"{_REPO}/resources/DeepCAD" +sys.path.insert(0, _DEEPCAD) +sys.path.insert(0, f"{_REPO}/resources/ll_gen_proof") +sys.path.insert(0, f"{_REPO}/ll_ocadr/mlx") + +from faithful_tower_mlx import FaithfulTower, precompute_geom # noqa: E402 + +from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh # noqa: E402 +from OCC.Core.TopExp import TopExp_Explorer # noqa: E402 +from OCC.Core.TopAbs import TopAbs_FACE # noqa: E402 +from OCC.Core.TopoDS import topods # noqa: E402 +from OCC.Core.BRep import BRep_Tool # noqa: E402 +from OCC.Core.TopLoc import TopLoc_Location # noqa: E402 + +NUM_POINTS, NUM_MESH_TOKENS = 2048, 256 +CLASS_NAMES = ["simple", "box", "complex"] + + +def _bucket(nf: int) -> int: + return 0 if nf <= 4 else (1 if nf <= 6 else 2) + + +def _tri_nodes(tri, k): + t = tri.Triangle(k) + try: + a, b, c = t.Get() + return a, b, c + except Exception: + return t.Value(1), t.Value(2), t.Value(3) + + +def _sample_points(shape, N=NUM_POINTS): + """Sample N (coords, normals) from the triangulated solid. Normals are + area-weighted vertex normals — the real encoders consume xyz + normals (6 ch).""" + BRepMesh_IncrementalMesh(shape, 0.05) + verts, tris = [], [] + exp = TopExp_Explorer(shape, TopAbs_FACE) + nfaces = 0 + while exp.More(): + nfaces += 1 + face = topods.Face(exp.Current()) + loc = TopLoc_Location() + tri = BRep_Tool.Triangulation(face, loc) + if tri is not None: + tr = loc.Transformation() + base = len(verts) + for i in range(1, tri.NbNodes() + 1): + p = tri.Node(i).Transformed(tr) + verts.append([p.X(), p.Y(), p.Z()]) + for k in range(1, tri.NbTriangles() + 1): + a, b, c = _tri_nodes(tri, k) + tris.append([base + a - 1, base + b - 1, base + c - 1]) + exp.Next() + if len(verts) < 3 or not tris: + return None, None, 0 + verts = np.asarray(verts, np.float32) + tris = np.asarray(tris, np.int64) + nrm = np.zeros_like(verts) + v0, v1, v2 = verts[tris[:, 0]], verts[tris[:, 1]], verts[tris[:, 2]] + fn = np.cross(v1 - v0, v2 - v0) # area-weighted face normals + for j in range(3): + np.add.at(nrm, tris[:, j], fn) + ln = np.linalg.norm(nrm, axis=1, keepdims=True) + nrm = np.where(ln > 1e-8, nrm / (ln + 1e-8), np.array([0.0, 0.0, 1.0], np.float32)) + c = verts.mean(0) + s = np.abs(verts - c).max() + 1e-6 + verts = (verts - c) / s + idx = np.random.choice(len(verts), N, replace=len(verts) < N) + return verts[idx].astype(np.float32), nrm[idx].astype(np.float32), nfaces + + +def build_dataset(n_target, deps, cache): + if cache and os.path.exists(cache): + d = np.load(cache) + if int(d["classes"].shape[0]) >= n_target: + return {k: d[k][:n_target] for k in + ("coords", "normals", "sa1", "sa2xyz", "sa2idx", "classes")} + CADSequence, command_dicts, execute = deps + files = sorted(glob.glob(os.path.join(_DEEPCAD, "data/cad_vec/*/*.h5"))) + coords, normals, sa1, sa2xyz, sa2idx, classes = [], [], [], [], [], [] + for f in files: + if len(classes) >= n_target: + break + try: + with h5py.File(f, "r") as h: + vec = h["vec"][:].astype(int) + cad = CADSequence.from_vector(vec, is_numerical=True, n=256) + shape = execute(command_dicts(cad)) + if shape is None: + continue + pc, nrm, nf = _sample_points(shape) + if pc is None or nf < 1: + continue + pre = precompute_geom(pc, nrm) + coords.append(pc) + normals.append(nrm) + sa1.append(pre["sa1_grouped"]) + sa2xyz.append(pre["sa2_grouped_xyz"]) + sa2idx.append(pre["sa2_group_idx"]) + classes.append(_bucket(nf)) + if len(classes) % 200 == 0: + print(f" built {len(classes)}/{n_target}", flush=True) + except Exception: + continue + out = {"coords": np.stack(coords), "normals": np.stack(normals), + "sa1": np.stack(sa1), "sa2xyz": np.stack(sa2xyz), + "sa2idx": np.stack(sa2idx), "classes": np.array(classes)} + if cache: + np.savez(cache, **out) + return out + + +class OCADRMLX(mlxnn.Module): + """Trainable faithful tower + LoRA-adapted (otherwise frozen) LLM, so mlx + value_and_grad differentiates tower params AND LoRA params together.""" + + def __init__(self, tower, llm): + super().__init__() + self.tower = tower + self.llm = llm + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "train"], default="train") + ap.add_argument("--llm", default="mlx-community/Qwen2-0.5B-Instruct-4bit") + ap.add_argument("--n-train", type=int, default=2200) + ap.add_argument("--n-val", type=int, default=400) + ap.add_argument("--epochs", type=int, default=10) + ap.add_argument("--bs", type=int, default=8) + ap.add_argument("--lr", type=float, default=2e-4) + ap.add_argument("--lora-layers", type=int, default=8) + ap.add_argument("--lora-rank", type=int, default=8) + ap.add_argument("--out", default=f"{_REPO}/ll_ocadr/checkpoints") + args = ap.parse_args() + + from cadlib.extrude import CADSequence + from ll_gen.proposals.command_proposal import CommandSequenceProposal + from ll_gen.disposal.command_executor import execute_command_proposal + import deepcad_supervised_train as M + + def command_dicts(cad): + cmds = M.translate(cad) + return CommandSequenceProposal(command_dicts=M.cad_to_command_dicts(cmds), + quantization_bits=8, normalization_range=2.0) + deps = (CADSequence, command_dicts, execute_command_proposal) + + print(f"loading {args.llm} + LoRA (rank={args.lora_rank}, layers={args.lora_layers}) ...", flush=True) + llm, tok = mlx_load(args.llm) + llm.freeze() + linear_to_lora_layers(llm, args.lora_layers, {"rank": args.lora_rank, "scale": 20.0, "dropout": 0.0}) + H = llm.args.hidden_size + tower = FaithfulTower(H, shape_depth=4, shape_heads=8) # 0.5B faithful config + model = OCADRMLX(tower, llm) + n_trainable = sum(v.size for _, v in tree_flatten(model.trainable_parameters())) + print(f"trainable (faithful tower + LoRA): {n_trainable/1e6:.2f}M; LLM base frozen+quantized", flush=True) + + prompt_ids = mx.array(tok.encode("Describe this CAD part: ")) + Lp = int(prompt_ids.shape[0]) + + def mesh_tokens(d, idx): + b = {"sa1_grouped": mx.array(d["sa1"][idx]), "sa2_grouped_xyz": mx.array(d["sa2xyz"][idx]), + "sa2_group_idx": mx.array(d["sa2idx"][idx].astype(np.int32)), + "coords": mx.array(d["coords"][idx]), "normals": mx.array(d["normals"][idx])} + return model.tower(b) # [B,256,H] + + def make_batch(d, idx, texts): + mesh = mesh_tokens(d, idx) + prompt_emb = model.llm.model.embed_tokens(prompt_ids) + resp = [mx.array(tok.encode(t + tok.eos_token)) for t in texts] + maxLr = max(int(r.shape[0]) for r in resp) + T = mesh.shape[1] + seqs, labels = [], [] + for b in range(len(texts)): + r = resp[b] + Lr = int(r.shape[0]) + remb = model.llm.model.embed_tokens(r) + emb = mx.concatenate([prompt_emb, mesh[b], remb], axis=0) + pad = maxLr - Lr + if pad: + emb = mx.concatenate([emb, mx.zeros((pad, H))], axis=0) + seqs.append(emb) + labels.append([-100] * (Lp + T) + r.tolist() + [-100] * pad) + return mx.stack(seqs), mx.array(np.array(labels)), mesh + + def loss_fn(d, idx, texts, classes): + emb, labels, mesh = make_batch(d, idx, texts) + logits = model.llm(mx.zeros(emb.shape[:2], dtype=mx.int32), input_embeddings=emb) + lg = logits[:, :-1].reshape(-1, logits.shape[-1]) + tg = labels[:, 1:].reshape(-1) + mask = tg != -100 + ce = mlxnn.losses.cross_entropy(lg, mx.where(mask, tg, 0), reduction="none") * mask + lm_loss = ce.sum() / mx.maximum(mask.sum(), 1) + aux = mlxnn.losses.cross_entropy(model.tower.aux_logits(mesh), mx.array(classes), reduction="mean") + return lm_loss + 0.5 * aux + + if args.mode == "probe": + d = build_dataset(6, deps, None) + model.tower.train() + lv, grads = mlxnn.value_and_grad(model, loss_fn)( + d, np.arange(4), [CLASS_NAMES[c] for c in d["classes"][:4]], d["classes"][:4]) + gnorm = float(mx.sqrt(sum((g * g).sum() for _, g in tree_flatten(grads))).item()) + mx.eval(lv) + print(f"probe loss={float(lv.item()):.4f} grad_norm={gnorm:.4f}", flush=True) + return + + os.makedirs(args.out, exist_ok=True) + cache = f"{args.out}/ocadr_faithful_data.npz" + print("building/loading dataset (N=2048 pts + normals + cached FPS grouping) ...", flush=True) + data = build_dataset(args.n_train + args.n_val, deps, cache) + allcl = data["classes"] + rng = np.random.default_rng(0) + per = int(np.bincount(allcl, minlength=3).min()) + keep = np.concatenate([rng.permutation(np.where(allcl == c)[0])[:per] for c in range(3)]) + rng.shuffle(keep) + data = {k: v[keep] for k, v in data.items()} + allcl = data["classes"] + nval = max(len(keep) // 6, 60) + val_idx = np.arange(nval) + tr_idx = np.arange(nval, len(keep)) + vcl, tcl = allcl[val_idx], allcl[tr_idx] + majority = float(np.bincount(vcl, minlength=3).max() / len(vcl)) + print(f"balanced {len(tr_idx)} train / {len(val_idx)} val; per-class={per}; " + f"val dist={np.bincount(vcl, minlength=3).tolist()} (majority={majority:.3f})", flush=True) + + opt = optim.Adam(learning_rate=args.lr) + lg_fn = mlxnn.value_and_grad(model, loss_fn) + + def greedy_class(d, idx): + B = len(idx) + model.tower.eval() + mesh = mesh_tokens(d, idx) + pe = mx.broadcast_to(model.llm.model.embed_tokens(prompt_ids)[None], (B, Lp, H)) + cur = mx.concatenate([pe, mesh], axis=1) + toks = [[] for _ in range(B)] + for _ in range(8): + logits = model.llm(mx.zeros(cur.shape[:2], dtype=mx.int32), input_embeddings=cur) + nxt = mx.argmax(logits[:, -1], axis=-1) + mx.eval(nxt) + cur = mx.concatenate([cur, model.llm.model.embed_tokens(nxt)[:, None]], axis=1) + for b in range(B): + toks[b].append(int(nxt[b].item())) + preds = np.full(B, -1, int) + for b in range(B): + txt = tok.decode(toks[b]).lower() + pos = {nm: txt.find(nm) for nm in CLASS_NAMES if nm in txt} + if pos: + preds[b] = CLASS_NAMES.index(min(pos, key=pos.get)) + return preds + + def class_acc(d, idx, classes, shuffle=False): + mesh_idx = idx[np.random.permutation(len(idx))] if shuffle else idx + correct = 0 + for k in range(0, len(idx), 16): + sl = slice(k, k + 16) + correct += int((greedy_class(d, mesh_idx[sl]) == classes[sl]).sum()) + return correct / len(idx) + + def aux_acc(d, idx, classes): + model.tower.eval() + correct = 0 + for k in range(0, len(idx), 16): + sl = slice(k, k + 16) + mesh = mesh_tokens(d, idx[sl]) + pred = np.array(mx.argmax(model.tower.aux_logits(mesh), axis=1).tolist()) + correct += int((pred == classes[sl]).sum()) + return correct / len(idx) + + best = -1.0 + n = len(tr_idx) + for epoch in range(args.epochs): + model.tower.train() + perm = np.random.permutation(n) + tot = 0.0 + nb = 0 + for k in range(0, n, args.bs): + bidx = tr_idx[perm[k:k + args.bs]] + lv, grads = lg_fn(data, bidx, [CLASS_NAMES[c] for c in data["classes"][bidx]], + data["classes"][bidx]) + opt.update(model, grads) + mx.eval(model.trainable_parameters(), opt.state, lv) + tot += float(lv.item()) + nb += 1 + acc = class_acc(data, val_idx, vcl) + shuf = class_acc(data, val_idx, vcl, shuffle=True) + aux = aux_acc(data, val_idx, vcl) + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} " + f"llm_gen_acc={acc:.3f} shuffled={shuf:.3f} encoder_aux_acc={aux:.3f}", flush=True) + if acc > best: + best = acc + mx.save_safetensors(f"{args.out}/ocadr_mlx.safetensors", + dict(tree_flatten(model.trainable_parameters()))) + + result = {"framework": "MLX (Apple Silicon)", "port": "faithful PointNet++/Point-BERT tower (parity-proven), retrained", + "task": "CAD point-cloud -> class (simple/box/complex), 3-way balanced", + "llm": args.llm, "llm_base": "frozen + 4-bit quantized", + "trainable": "faithful GeometryNet+ShapeNet+projector + LoRA", "trainable_params_M": round(n_trainable / 1e6, 2), + "n_train": int(len(tr_idx)), "n_val": int(len(val_idx)), "epochs": args.epochs, + "num_mesh_tokens": NUM_MESH_TOKENS, "majority_baseline": round(majority, 3), + "encoder_mesh_read_acc": round(aux_acc(data, val_idx, vcl), 3), + "llm_generation_acc": round(class_acc(data, val_idx, vcl), 3), + "shuffled_mesh_baseline": round(class_acc(data, val_idx, vcl, shuffle=True), 3), + "checkpoint": f"{args.out}/ocadr_mlx.safetensors"} + with open(f"{args.out}/ocadr_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("OCADR_MLX_FAITHFUL_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_ocadr/tests/test_vision_modality.py b/ll_ocadr/tests/test_vision_modality.py index e28addc..48e6c23 100644 --- a/ll_ocadr/tests/test_vision_modality.py +++ b/ll_ocadr/tests/test_vision_modality.py @@ -121,6 +121,11 @@ class _StubLLM(torch.nn.Module): def __init__(self): super().__init__() self.embed = torch.nn.Embedding(50, n_embed) + # Real HF models expose .config.hidden_size; the model derives + # n_embed from it so the projector/splice match the LLM width. + self.config = types.SimpleNamespace( + hidden_size=n_embed, vocab_size=50 + ) def get_input_embeddings(self): return self.embed diff --git a/ll_ocadr/tests/test_wiring.py b/ll_ocadr/tests/test_wiring.py new file mode 100644 index 0000000..0d9eaf9 --- /dev/null +++ b/ll_ocadr/tests/test_wiring.py @@ -0,0 +1,88 @@ +"""Regression: the LatticelabsOCADR multimodal wiring must run end to end. + +The configured model could not run before three fixes: + +1. ``config.n_embed`` was hardcoded (1280) and matched no real base LLM + (Qwen2-0.5B=896, 1.5B=1536, 7B=3584), so the projected mesh tokens (n_embed) + could not be spliced into the LLM input embeddings (hidden_size) — shape + mismatch. ``n_embed`` is now derived from the loaded LLM's hidden size. +2. ``get_config_for_model`` referenced non-existent Qwen2 sizes; it now uses + real sizes ("0.5b"/"1.5b"/"7b"). +3. ``_splice_tokens`` assigned fp32 modality embeddings into a bf16 LLM + embedding tensor — ``Index put requires ... dtypes match``. It now casts to + the destination dtype/device. + +These tests mock the base LLM so they do not download Qwen, exercising the +ll_ocadr-side wiring (encoders -> projector -> splice -> LLM) directly. +""" +from __future__ import annotations + +import types +from unittest.mock import patch + +import pytest + +torch = pytest.importorskip("torch") +import torch.nn as nn # noqa: E402 + + +class _FakeLLM(nn.Module): + """Stand-in base LLM with a known hidden size, loaded in bfloat16 (as real + Qwen2 checkpoints are), exposing the get_input_embeddings/forward surface + LatticelabsOCADR depends on.""" + + def __init__(self, hidden: int = 64, vocab: int = 128) -> None: + super().__init__() + self.config = types.SimpleNamespace(hidden_size=hidden, vocab_size=vocab) + self._embed = nn.Embedding(vocab, hidden) + self._head = nn.Linear(hidden, vocab) + self.to(torch.bfloat16) # emulate a bf16-loaded LLM + + def get_input_embeddings(self): + return self._embed + + def forward(self, inputs_embeds=None, attention_mask=None, **kwargs): + return types.SimpleNamespace(logits=self._head(inputs_embeds)) + + +def _build_model(hidden: int = 64, vocab: int = 128): + from ll_ocadr.vllm.config import get_config_for_model + from ll_ocadr.vllm import latticelabs_ocadr as M + + cfg = get_config_for_model("0.5b") + cfg.mesh_token_id = 7 + with patch.object(M, "AutoModelForCausalLM") as auto: + auto.from_pretrained.return_value = _FakeLLM(hidden, vocab) + model = M.LatticelabsOCADRForCausalLM(cfg).eval() + return cfg, model + + +def test_n_embed_is_derived_from_llm(): + """n_embed must equal the LLM hidden size, not the hardcoded 1280.""" + cfg, model = _build_model(hidden=64) + assert cfg.n_embed == 64 + assert model.language_model.config.hidden_size == 64 + # the projector output layer must map into the LLM hidden dim + assert model.projector(torch.randn(1, 3, cfg.input_dim)).shape[-1] == 64 + + +def test_forward_runs_and_splices_mesh_across_dtypes(): + """Full forward: mesh -> encoders -> projector -> bf16 splice -> LLM logits.""" + cfg, model = _build_model(hidden=64, vocab=128) + vc = torch.randn(1, 256, 3) + vn = torch.randn(1, 256, 3) + mesh = model._mesh_to_embedding(vc, vn) + n = mesh[0].shape[0] + assert mesh[0].shape[-1] == 64 # projected to LLM hidden dim + + input_ids = torch.cat( + [torch.tensor([[1, 2]]), torch.full((1, n), 7), torch.tensor([[3]])], dim=1 + ) + out = model.forward( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + vertex_coords=vc, + vertex_normals=vn, + ) + assert out.logits.shape == (1, input_ids.shape[1], 128) + assert torch.isfinite(out.logits.float()).all() diff --git a/ll_ocadr/vllm/config.py b/ll_ocadr/vllm/config.py index fe5cff2..035ec09 100644 --- a/ll_ocadr/vllm/config.py +++ b/ll_ocadr/vllm/config.py @@ -80,34 +80,42 @@ def get_default_config() -> LLOCADRConfig: def get_config_for_model(model_size: str = "7b") -> LLOCADRConfig: """ - Get configuration for specific model size. + Get configuration for a specific base-LLM size. + + ``n_embed`` (the LLM embedding/hidden dimension that the projector and all + learnable separators must match) is NOT set here — it is derived from the + actual loaded language model in ``LatticelabsOCADRForCausalLM.__init__``, so + every size wires correctly regardless of this table. Only the base LLM name + and the 3D shape-encoder capacity vary by size. Args: - model_size: "1.8b", "7b", or "14b" + model_size: "0.5b", "1.5b", or "7b" (real Qwen2 sizes that exist on the + Hub). "0.5b"/"1.5b" run/train on a single consumer machine; "7b" + needs a GPU. Returns: LLOCADRConfig """ base_config = LLOCADRConfig() - if model_size == "1.8b": - base_config.language_model_name = "Qwen/Qwen2-1.8B" - base_config.n_embed = 1024 + if model_size == "0.5b": + base_config.language_model_name = "Qwen/Qwen2-0.5B" + base_config.shape_depth = 4 + base_config.shape_num_heads = 8 + elif model_size == "1.5b": + base_config.language_model_name = "Qwen/Qwen2-1.5B" base_config.shape_depth = 6 base_config.shape_num_heads = 8 elif model_size == "7b": base_config.language_model_name = "Qwen/Qwen2-7B" - base_config.n_embed = 1280 base_config.shape_depth = 12 base_config.shape_num_heads = 12 - elif model_size == "14b": - base_config.language_model_name = "Qwen/Qwen2-14B" - base_config.n_embed = 1536 - base_config.shape_depth = 16 - base_config.shape_num_heads = 16 else: - raise ValueError(f"Unknown model size: {model_size}") + raise ValueError( + f"Unknown model size: {model_size!r} (expected '0.5b', '1.5b', '7b')" + ) + base_config.model_name = f"latticelabs/ll-ocadr-{model_size}" base_config.input_dim = base_config.geometry_embed_dim + base_config.shape_embed_dim return base_config diff --git a/ll_ocadr/vllm/latticelabs_ocadr.py b/ll_ocadr/vllm/latticelabs_ocadr.py index 2dc549c..bb114f7 100644 --- a/ll_ocadr/vllm/latticelabs_ocadr.py +++ b/ll_ocadr/vllm/latticelabs_ocadr.py @@ -22,6 +22,19 @@ def __init__(self, config): super().__init__() self.config = config + # Language model — loaded FIRST so every projection is sized to its + # hidden dimension. Mesh/image tokens are spliced INTO the LLM's input + # embeddings (see get_input_embeddings/_splice_tokens), so the projector + # output and all learnable separators MUST match the LLM hidden size. + # A hardcoded config.n_embed (1280) matches no real base + # (Qwen2-0.5B=896, 1.5B=1536, 7B=3584), so the splice raised a shape + # mismatch and the configured model could not run. Derive n_embed from + # the actual LLM so any base size wires correctly. + self.language_model = AutoModelForCausalLM.from_pretrained( + config.language_model_name + ) + config.n_embed = int(self.language_model.config.hidden_size) + # Initialize 3D encoders self.geometry_model = build_geometry_net() # Local geometry features self.shape_model = build_shape_net( @@ -30,14 +43,10 @@ def __init__(self, config): num_heads=config.shape_num_heads, ) # Global shape features - # MLP Projector: concatenated features -> LLM embedding space + # MLP Projector: concatenated features -> LLM embedding space (n_embed + # == LLM hidden size, set above). self.projector = MlpProjector(config) - # Language model - self.language_model = AutoModelForCausalLM.from_pretrained( - config.language_model_name - ) - # Special tokens (learnable parameters) self.mesh_boundary = nn.Parameter( torch.randn(1, config.n_embed) @@ -312,7 +321,14 @@ def _splice_tokens( .squeeze(-1) ) if len(positions) > 0 and batch_idx < len(embeddings): - emb = embeddings[batch_idx] + # The language model may run in reduced precision (its weights + # often load as bfloat16 while the 3D encoders/projector emit + # float32) and may sit on a different device. Cast the modality + # embeddings to the destination embeddings' dtype and device so + # the in-place splice does not raise a dtype/device mismatch. + emb = embeddings[batch_idx].to( + dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) num = emb.shape[0] if len(positions) >= num: inputs_embeds[batch_idx, positions[:num]] = emb diff --git a/ll_stepnet/mlx/train_classification_mlx.py b/ll_stepnet/mlx/train_classification_mlx.py new file mode 100644 index 0000000..5b14dfd --- /dev/null +++ b/ll_stepnet/mlx/train_classification_mlx.py @@ -0,0 +1,442 @@ +"""ll_stepnet STEPForClassification in native MLX — faithful weight-conversion port. + +This is NOT a simplified re-implementation. It reproduces the EXACT architecture of +``stepnet.tasks.STEPForClassification`` (the model that reached PyTorch val acc 0.976) +and CONVERTS the real trained PyTorch checkpoint +(``ll_stepnet/checkpoints/stepnet_classifier.pt``) into MLX, so the MLX model *is* the +trained model — same weights, same accuracy — running natively on Apple Silicon. + +Architecture (reconstructed from the checkpoint state_dict, see encoder.py / tasks.py): + + STEPForClassification + encoder (STEPEncoder) + transformer_encoder (STEPTransformerEncoder) + token_embedding : Embedding(50000, 256) + pos_embedding : (1, 5000, 256) # trained nn.Parameter + transformer : 6 x nn.TransformerEncoderLayer (post-norm, relu, 8 heads, ff=1024) + layer_norm : LayerNorm(256) # applied after the stack + graph_encoder (STEPGraphEncoder) # UNUSED at inference: topology=None -> zeros(128) + fusion : Linear(384->1024) -> ReLU -> Linear(1024->1024) + classifier : Linear(1024->512) -> ReLU -> Dropout -> Linear(512->3) + + forward(token_ids): # topology_data is None throughout training/eval + x = transformer_encoder(token_ids) # [B, S, 256] + token_pooled = x.mean(dim=1) # [B, 256] (UNMASKED mean over all positions) + combined = cat([token_pooled, zeros(B,128)]) # [B, 384] + return classifier(fusion(combined)) # [B, 3] + +Every checkpoint tensor maps 1:1 onto an MLX Linear/Embedding/LayerNorm of identical +shape (MLX Linear.weight is [out, in] just like PyTorch), so conversion is a direct +array copy with NO transposes — the only non-trivial piece is multi-head attention, +implemented here manually from the packed ``in_proj_weight`` (768,256) = [Wq;Wk;Wv] +so it is bit-faithful to PyTorch's F.multi_head_attention_forward. + +Modes: + probe - build the MLX model, run a forward, print shapes. + convert - load the real .pt, convert weights, save MLX safetensors + parity report. + parity - convert AND run BOTH the PyTorch and MLX models on the same real DeepCAD + val split; report argmax-agreement rate + each model's val accuracy. + train - train the faithful architecture from scratch in MLX (proves it is also + natively trainable, not just a frozen import). +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +import warnings +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MPLBACKEND", "Agg") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +matplotlib.use = lambda *a, **k: None # neutralize cadlib's TkAgg switch + +import numpy as np # noqa: E402 +import h5py # noqa: E402 +import mlx.core as mx # noqa: E402 +import mlx.nn as mlxnn # noqa: E402 +import mlx.optimizers as optim # noqa: E402 +from mlx.utils import tree_flatten, tree_unflatten # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] +_DEEPCAD = str(_REPO / "resources/DeepCAD") + +# --- DeepCAD cad_vec -> command-token sequence + face-count class -------------- +# IDENTICAL tokenization to ll_stepnet/scripts/train_classification.py (the trainer +# that produced the checkpoint), so the val data fed to the MLX model matches exactly. +LEVELS, RANGE, MAX_LEN = 256, 2.0, 256 +NUM_SLOTS = 16 +MASK = {"LINE": [0, 1, 2, 3], "ARC": [0, 1, 2, 3, 4, 5], "CIRCLE": [0, 1, 2], + "EXTRUDE": [0, 1, 2, 3, 4, 5, 6, 7], "SOL": [], "EOS": []} +CMD_TOK = {"SOL": 6, "LINE": 7, "ARC": 8, "CIRCLE": 9, "EXTRUDE": 10, "EOS": 11} +BUCKETS = [(0, 4), (5, 6), (7, 9999)] +CLASS_NAMES = ["simple(<=4)", "box(5-6)", "complex(7+)"] +VOCAB = 50000 # matches the checkpoint's token_embedding (50000, 256) + + +def _q_coord(g): + return int(np.clip(round(float(g)), 0, LEVELS - 1)) + + +def _q_value(v): + return int(np.clip(round((float(v) + RANGE) / (2 * RANGE) * (LEVELS - 1)), 0, LEVELS - 1)) + + +def _translate(cad, Circle, Arc): + cmds = [] + for ext in cad.seq: + for loop in ext.profile.children: + cmds.append(("SOL", {})) + for cv in loop.children: + if isinstance(cv, Circle): + r_mag = float(cv.radius) / (LEVELS - 1) * 2.0 * RANGE + cmds.append(("CIRCLE", {0: _q_coord(cv.center[0]), 1: _q_coord(cv.center[1]), + 2: _q_value(r_mag)})) + elif isinstance(cv, Arc): + s, e, c = cv.start_point, cv.end_point, cv.center + cmds.append(("ARC", {0: _q_coord(s[0]), 1: _q_coord(s[1]), 2: _q_coord(e[0]), + 3: _q_coord(e[1]), 4: _q_coord(c[0]), 5: _q_coord(c[1])})) + else: + s, e = cv.start_point, cv.end_point + cmds.append(("LINE", {0: _q_coord(s[0]), 1: _q_coord(s[1]), + 2: _q_coord(e[0]), 3: _q_coord(e[1])})) + depth = abs(float(ext.extent_one)) + abs(float(ext.extent_two)) + cmds.append(("EXTRUDE", {0: _q_value(float(np.clip(depth * 4.0, 0.3, 2.0)))})) + cmds.append(("EOS", {})) + return cmds + + +def _command_dicts(cmds): + out = [] + for name, slots in cmds: + p = [0] * NUM_SLOTS + m = [False] * NUM_SLOTS + for j in MASK[name]: + p[j] = int(slots.get(j, 0)) + m[j] = True + out.append({"command_type": name, "parameters": p, "parameter_mask": m}) + return out + + +def _encode_tokens(cmds): + toks = [1] + for name, slots in cmds: + toks.append(CMD_TOK[name]) + for j in MASK[name]: + toks.append(12 + int(slots.get(j, 0))) + toks.append(2) + toks = toks[:MAX_LEN] + return toks + [0] * (MAX_LEN - len(toks)) + + +def _bucket(nf): + for i, (lo, hi) in enumerate(BUCKETS): + if lo <= nf <= hi: + return i + return len(BUCKETS) - 1 + + +def build_val_split(n_val, cache): + """Reproduce the PyTorch trainer's val split EXACTLY: scan sorted(files)[:need//6] + and collect the first ``n_val`` that produce valid solids, same tokenization.""" + if cache and os.path.exists(cache): + d = np.load(cache) + if d["tokens"].shape[0] >= n_val: + return d["tokens"][:n_val], d["classes"][:n_val] + sys.path.insert(0, _DEEPCAD) + sys.path.insert(0, str(_REPO / "resources/ll_gen_proof")) + from cadlib.extrude import CADSequence + from cadlib.curves import Arc, Circle + from ll_gen.proposals.command_proposal import CommandSequenceProposal + from ll_gen.disposal.command_executor import execute_command_proposal + from OCC.Core.TopExp import TopExp_Explorer + from OCC.Core.TopAbs import TopAbs_FACE + + def face_count(shape): + c = 0 + e = TopExp_Explorer(shape, TopAbs_FACE) + while e.More(): + c += 1 + e.Next() + return c + + files = sorted(glob.glob(os.path.join(_DEEPCAD, "data/cad_vec", "*/*.h5"))) + need = (5000 + 1000) * 3 + 4000 # == 22000, mirrors the trainer + files = files[: need // 6] + toks, classes = [], [] + for f in files: + if len(toks) >= n_val: + break + try: + with h5py.File(f, "r") as h: + vec = h["vec"][:].astype(int) + cad = CADSequence.from_vector(vec, is_numerical=True, n=256) + cmds = _translate(cad, Circle, Arc) + shape = execute_command_proposal(CommandSequenceProposal( + command_dicts=_command_dicts(cmds), quantization_bits=8, normalization_range=2.0)) + if shape is None: + continue + nf = face_count(shape) + if nf < 1: + continue + toks.append(_encode_tokens(cmds)) + classes.append(_bucket(nf)) + except Exception: + continue + toks = np.array(toks, np.int32) + classes = np.array(classes, np.int32) + if cache: + np.savez(cache, tokens=toks, classes=classes) + return toks, classes + + +# --- faithful MLX model (exact port of STEPForClassification) ------------------ +class FaithfulMHA(mlxnn.Module): + """PyTorch nn.MultiheadAttention math from the packed in_proj_weight (768,256).""" + + def __init__(self, d=256, heads=8): + super().__init__() + self.d, self.h, self.hd = d, heads, d // heads + self.in_proj = mlxnn.Linear(d, 3 * d) # weight (768,256), bias (768,) + self.out_proj = mlxnn.Linear(d, d) # weight (256,256), bias (256,) + + def __call__(self, x): # x: [B, S, d] + B, S, _ = x.shape + qkv = self.in_proj(x) # [B,S,768] + q, k, v = mx.split(qkv, 3, axis=-1) # 3 x [B,S,256] + + def heads(t): + return mx.transpose(t.reshape(B, S, self.h, self.hd), (0, 2, 1, 3)) # [B,h,S,hd] + + q, k, v = heads(q), heads(k), heads(v) + scores = (q @ mx.transpose(k, (0, 1, 3, 2))) / (self.hd ** 0.5) # [B,h,S,S] + ctx = mx.softmax(scores, axis=-1) @ v # [B,h,S,hd] + ctx = mx.transpose(ctx, (0, 2, 1, 3)).reshape(B, S, self.d) # [B,S,d] + return self.out_proj(ctx) + + +class FaithfulLayer(mlxnn.Module): + """nn.TransformerEncoderLayer (post-norm, activation=relu) — PyTorch defaults.""" + + def __init__(self, d=256, heads=8, ff=1024): + super().__init__() + self.self_attn = FaithfulMHA(d, heads) + self.linear1 = mlxnn.Linear(d, ff) + self.linear2 = mlxnn.Linear(ff, d) + self.norm1 = mlxnn.LayerNorm(d) + self.norm2 = mlxnn.LayerNorm(d) + + def __call__(self, x): + x = self.norm1(x + self.self_attn(x)) + return self.norm2(x + self.linear2(mlxnn.relu(self.linear1(x)))) + + +class FaithfulTransformerEncoder(mlxnn.Module): + def __init__(self, vocab=VOCAB, d=256, layers=6, heads=8, ff=1024): + super().__init__() + self.token_embedding = mlxnn.Embedding(vocab, d) + self.pos_embedding = mx.zeros((1, 5000, d)) # overwritten by converted weights + self.layers = [FaithfulLayer(d, heads, ff) for _ in range(layers)] + self.layer_norm = mlxnn.LayerNorm(d) + + def __call__(self, ids): + x = self.token_embedding(ids) + self.pos_embedding[:, : ids.shape[1], :] + for layer in self.layers: + x = layer(x) + return self.layer_norm(x) + + +class FaithfulSTEPClassifier(mlxnn.Module): + def __init__(self, vocab=VOCAB, nclass=3, d=256, output_dim=1024, graph_node_dim=128): + super().__init__() + self.te = FaithfulTransformerEncoder(vocab, d) + self.graph_node_dim = graph_node_dim + self.fusion0 = mlxnn.Linear(d + graph_node_dim, output_dim) + self.fusion2 = mlxnn.Linear(output_dim, output_dim) + self.cls0 = mlxnn.Linear(output_dim, 512) + self.cls3 = mlxnn.Linear(512, nclass) + + def __call__(self, ids): + x = self.te(ids) # [B,S,256] + token_pooled = x.mean(axis=1) # [B,256] unmasked mean + zero_graph = mx.zeros((ids.shape[0], self.graph_node_dim)) + combined = mx.concatenate([token_pooled, zero_graph], axis=-1) # [B,384] + out = self.fusion2(mlxnn.relu(self.fusion0(combined))) # [B,1024] + return self.cls3(mlxnn.relu(self.cls0(out))) # [B,nclass] + + +# --- weight conversion: real PyTorch checkpoint -> MLX params ------------------ +def convert_checkpoint(ckpt_path, model): + """Load the real trained state_dict and assign it into the MLX model by an + explicit name map. Returns (model, n_converted).""" + import torch + + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + sd = ck.get("model_state_dict", ck) + + def arr(key): + return mx.array(sd[key].detach().cpu().float().numpy()) + + pairs = [] + te = "encoder.transformer_encoder" + pairs.append(("te.token_embedding.weight", arr(f"{te}.token_embedding.weight"))) + pairs.append(("te.pos_embedding", arr(f"{te}.pos_embedding"))) + pairs.append(("te.layer_norm.weight", arr(f"{te}.layer_norm.weight"))) + pairs.append(("te.layer_norm.bias", arr(f"{te}.layer_norm.bias"))) + n_layers = len({k.split(".layers.")[1].split(".")[0] + for k in sd if f"{te}.transformer.layers." in k}) + for i in range(n_layers): + s = f"{te}.transformer.layers.{i}" + d = f"te.layers.{i}" + pairs += [ + (f"{d}.self_attn.in_proj.weight", arr(f"{s}.self_attn.in_proj_weight")), + (f"{d}.self_attn.in_proj.bias", arr(f"{s}.self_attn.in_proj_bias")), + (f"{d}.self_attn.out_proj.weight", arr(f"{s}.self_attn.out_proj.weight")), + (f"{d}.self_attn.out_proj.bias", arr(f"{s}.self_attn.out_proj.bias")), + (f"{d}.linear1.weight", arr(f"{s}.linear1.weight")), + (f"{d}.linear1.bias", arr(f"{s}.linear1.bias")), + (f"{d}.linear2.weight", arr(f"{s}.linear2.weight")), + (f"{d}.linear2.bias", arr(f"{s}.linear2.bias")), + (f"{d}.norm1.weight", arr(f"{s}.norm1.weight")), + (f"{d}.norm1.bias", arr(f"{s}.norm1.bias")), + (f"{d}.norm2.weight", arr(f"{s}.norm2.weight")), + (f"{d}.norm2.bias", arr(f"{s}.norm2.bias")), + ] + for src, dst in (("encoder.fusion.0", "fusion0"), ("encoder.fusion.2", "fusion2"), + ("classifier.0", "cls0"), ("classifier.3", "cls3")): + pairs.append((f"{dst}.weight", arr(f"{src}.weight"))) + pairs.append((f"{dst}.bias", arr(f"{src}.bias"))) + + model.update(tree_unflatten(pairs)) + mx.eval(model.parameters()) + return model, len(pairs) + + +def _pt_logits(ckpt_path, toks, nclass): + """Run the real PyTorch STEPForClassification on the same tokens (for parity).""" + import torch + sys.path.insert(0, str(_REPO / "ll_stepnet")) + from stepnet.tasks import STEPForClassification + + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + model = STEPForClassification(num_classes=nclass) + model.load_state_dict(ck["model_state_dict"], strict=False) + model.eval() + outs = [] + with torch.no_grad(): + for k in range(0, toks.shape[0], 256): + t = torch.tensor(toks[k:k + 256], dtype=torch.long) + outs.append(model(t).cpu().numpy()) + return np.concatenate(outs, axis=0) + + +def _mlx_logits(model, toks): + outs = [] + for k in range(0, toks.shape[0], 256): + outs.append(np.array(model(mx.array(toks[k:k + 256])).tolist())) + return np.concatenate(outs, axis=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "convert", "parity", "train"], default="parity") + ap.add_argument("--ckpt", default=str(_REPO / "ll_stepnet/checkpoints/stepnet_classifier.pt")) + ap.add_argument("--n-val", type=int, default=1000) + ap.add_argument("--epochs", type=int, default=30) + ap.add_argument("--bs", type=int, default=64) + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--out", default=str(_REPO / "ll_stepnet/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + nclass = len(BUCKETS) + + if args.mode == "probe": + model = FaithfulSTEPClassifier(nclass=nclass) + ids = mx.array(np.random.randint(0, 268, (4, MAX_LEN)).astype(np.int32)) + out = model(ids) + print(f"probe: logits {out.shape} finite={bool(mx.isfinite(out).all().item())}", flush=True) + return + + if args.mode == "convert": + model = FaithfulSTEPClassifier(nclass=nclass) + model, n = convert_checkpoint(args.ckpt, model) + out_path = f"{args.out}/stepnet_classifier_mlx.safetensors" + mx.save_safetensors(out_path, dict(tree_flatten(model.parameters()))) + print(f"converted {n} tensors -> {out_path}", flush=True) + return + + if args.mode == "parity": + cache = f"{args.out}/mlx_parity_val.npz" + print(f"building/loading {args.n_val} real val samples (same split as trainer) ...", flush=True) + toks, cls = build_val_split(args.n_val, cache) + print(f"val set: {toks.shape[0]} samples, dist={np.bincount(cls, minlength=nclass).tolist()}", flush=True) + + model = FaithfulSTEPClassifier(nclass=nclass) + model, n = convert_checkpoint(args.ckpt, model) + mx.save_safetensors(f"{args.out}/stepnet_classifier_mlx.safetensors", + dict(tree_flatten(model.parameters()))) + print(f"converted {n} real tensors into MLX", flush=True) + + lg_mlx = _mlx_logits(model, toks) + lg_pt = _pt_logits(args.ckpt, toks, nclass) + pred_mlx, pred_pt = lg_mlx.argmax(1), lg_pt.argmax(1) + agree = float((pred_mlx == pred_pt).mean()) + max_logit_diff = float(np.abs(lg_mlx - lg_pt).max()) + acc_mlx = float((pred_mlx == cls).mean()) + acc_pt = float((pred_pt == cls).mean()) + per = {CLASS_NAMES[c]: round(float(((pred_mlx == cls) & (cls == c)).sum() / + max((cls == c).sum(), 1)), 3) for c in range(nclass)} + majority = float(np.bincount(cls, minlength=nclass).max() / len(cls)) + + result = {"framework": "MLX (Apple Silicon)", "port": "faithful weight-conversion", + "task": "STEP->face-count complexity class (3)", "dataset": "DeepCAD cad_vec", + "n_val": int(toks.shape[0]), "source_checkpoint": args.ckpt, + "argmax_agreement_vs_pytorch": round(agree, 4), + "max_abs_logit_diff": round(max_logit_diff, 5), + "mlx_val_acc": round(acc_mlx, 4), "pytorch_val_acc": round(acc_pt, 4), + "majority_baseline": round(majority, 3), "mlx_per_class_acc": per, + "checkpoint": f"{args.out}/stepnet_classifier_mlx.safetensors"} + with open(f"{args.out}/stepnet_classifier_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("STEPNET_MLX_PARITY", json.dumps(result), flush=True) + return + + # mode == train : native MLX training of the faithful architecture from scratch + cache = f"{args.out}/mlx_parity_val.npz" + toks, cls = build_val_split(args.n_val, cache) + n_tr = int(toks.shape[0] * 0.8) + tt, tcl, vt, vcl = toks[:n_tr], cls[:n_tr], toks[n_tr:], cls[n_tr:] + model = FaithfulSTEPClassifier(nclass=nclass) + cnt = np.bincount(tcl, minlength=nclass) + w = mx.array((cnt.sum() / (nclass * np.clip(cnt, 1, None))).astype(np.float32)) + opt = optim.Adam(learning_rate=args.lr) + + def loss_fn(ids, y): + ce = mlxnn.losses.cross_entropy(model(ids), y, reduction="none") + return (ce * w[y]).mean() + + lg = mlxnn.value_and_grad(model, loss_fn) + for epoch in range(args.epochs): + perm = np.random.permutation(tt.shape[0]) + for k in range(0, tt.shape[0], args.bs): + idx = perm[k:k + args.bs] + lv, g = lg(mx.array(tt[idx]), mx.array(tcl[idx])) + opt.update(model, g) + mx.eval(model.parameters(), opt.state, lv) + pv = _mlx_logits(model, vt).argmax(1) + print(f"epoch {epoch+1}/{args.epochs} val_acc={float((pv == vcl).mean()):.3f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_stepnet/scripts/train_classification.py b/ll_stepnet/scripts/train_classification.py new file mode 100644 index 0000000..62e453f --- /dev/null +++ b/ll_stepnet/scripts/train_classification.py @@ -0,0 +1,262 @@ +"""Train ll_stepnet's STEPForClassification on real DeepCAD models. + +First real trained ll_stepnet checkpoint (make-real campaign). Task: predict a +model's face-count complexity class (<=4 / 5-6 / 7+ faces) from its CAD command +sequence — a GEOMETRIC label derived from the reconstructed solid, so the encoder +must learn the construction->geometry relationship rather than count tokens. + +Pipeline: DeepCAD cad_vec (h5) -> cadlib CADSequence (absolute geometry) -> +executor-schema command_dicts -> OCC solid (for the face-count label) + flattened +token sequence (for the input). Trains stepnet.tasks.STEPForClassification and +writes the checkpoint + metrics to ll_stepnet/checkpoints/. + +Requires the canonical DeepCAD data extracted under ``--deepcad`` (default +``resources/DeepCAD/data/cad_vec``; download via +http://www.cs.columbia.edu/cg/deepcad/data.tar) and the ll_gen executor + +pythonocc (conda 'cadling' env). + +Run:: + + python ll_stepnet/scripts/train_classification.py --n-train 5000 --epochs 30 \ + --device mps +""" + +from __future__ import annotations + +import argparse +import collections +import glob +import json +import logging +import os +import sys +import warnings +from pathlib import Path + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MPLBACKEND", "Agg") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +warnings.filterwarnings("ignore") +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +matplotlib.use = lambda *a, **k: None # neutralize cadlib's TkAgg switch +logging.disable(logging.WARNING) + +import numpy as np # noqa: E402 +import h5py # noqa: E402 +import torch # noqa: E402 +import torch.nn as nn # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] + +# ---------------------------------------------------------------------------- +# DeepCAD cad_vec -> executor-schema translation (self-contained; validated +# 30/30 real models -> valid solids). Mirrors the executor's symmetric +# [0,255]<->[-2,2] param quantization. +# ---------------------------------------------------------------------------- +LEVELS, RANGE, MAX_CMDS, NUM_SLOTS = 256, 2.0, 60, 16 +MASK = {"LINE": [0, 1, 2, 3], "ARC": [0, 1, 2, 3, 4, 5], "CIRCLE": [0, 1, 2], + "EXTRUDE": [0, 1, 2, 3, 4, 5, 6, 7], "SOL": [], "EOS": []} +CMD_TOK = {"SOL": 6, "LINE": 7, "ARC": 8, "CIRCLE": 9, "EXTRUDE": 10, "EOS": 11} +MAX_LEN = 256 +BUCKETS = [(0, 4), (5, 6), (7, 9999)] +CLASS_NAMES = ["simple(<=4)", "box(5-6)", "complex(7+)"] + + +def _q_coord(grid: float) -> int: + return int(np.clip(round(float(grid)), 0, LEVELS - 1)) + + +def _q_value(v: float) -> int: + return int(np.clip(round((float(v) + RANGE) / (2 * RANGE) * (LEVELS - 1)), 0, LEVELS - 1)) + + +def _translate(cad, Circle, Arc): + """CADSequence (absolute) -> list of (command_name, {slot: quant}).""" + cmds = [] + for ext in cad.seq: + for loop in ext.profile.children: + cmds.append(("SOL", {})) + for cv in loop.children: + if isinstance(cv, Circle): + r_mag = float(cv.radius) / (LEVELS - 1) * 2.0 * RANGE + cmds.append(("CIRCLE", {0: _q_coord(cv.center[0]), 1: _q_coord(cv.center[1]), + 2: _q_value(r_mag)})) + elif isinstance(cv, Arc): + s, e, c = cv.start_point, cv.end_point, cv.center + cmds.append(("ARC", {0: _q_coord(s[0]), 1: _q_coord(s[1]), 2: _q_coord(e[0]), + 3: _q_coord(e[1]), 4: _q_coord(c[0]), 5: _q_coord(c[1])})) + else: + s, e = cv.start_point, cv.end_point + cmds.append(("LINE", {0: _q_coord(s[0]), 1: _q_coord(s[1]), + 2: _q_coord(e[0]), 3: _q_coord(e[1])})) + depth = abs(float(ext.extent_one)) + abs(float(ext.extent_two)) + cmds.append(("EXTRUDE", {0: _q_value(float(np.clip(depth * 4.0, 0.3, 2.0)))})) + cmds.append(("EOS", {})) + return cmds + + +def _command_dicts(cmds): + out = [] + for name, slots in cmds: + p = [0] * NUM_SLOTS + m = [False] * NUM_SLOTS + for j in MASK[name]: + p[j] = int(slots.get(j, 0)) + m[j] = True + out.append({"command_type": name, "parameters": p, "parameter_mask": m}) + return out + + +def _encode_tokens(cmds): + toks = [1] # BOS + for name, slots in cmds: + toks.append(CMD_TOK[name]) + for j in MASK[name]: + toks.append(12 + int(slots.get(j, 0))) + toks.append(2) # EOS + toks = toks[:MAX_LEN] + toks += [0] * (MAX_LEN - len(toks)) + return toks + + +def _bucket(nf): + for i, (lo, hi) in enumerate(BUCKETS): + if lo <= nf <= hi: + return i + return len(BUCKETS) - 1 + + +def load_dataset(files, limit, deps): + CADSequence, Circle, Arc, CommandSequenceProposal, execute, face_count = deps + toks, labels = [], [] + for f in files: + if len(toks) >= limit: + break + try: + with h5py.File(f, "r") as h: + vec = h["vec"][:].astype(int) + cad = CADSequence.from_vector(vec, is_numerical=True, n=256) + cmds = _translate(cad, Circle, Arc) + shape = execute(CommandSequenceProposal( + command_dicts=_command_dicts(cmds), quantization_bits=8, normalization_range=2.0)) + if shape is None: + continue + nf = face_count(shape) + if nf < 1: + continue + toks.append(_encode_tokens(cmds)) + labels.append(_bucket(nf)) + except Exception: + continue + return torch.tensor(toks, dtype=torch.long), torch.tensor(labels, dtype=torch.long) + + +def accuracy(model, tok, lab, dev, bs=256): + model.eval() + corr = tot = 0 + per = np.zeros((len(BUCKETS), 2), dtype=int) + with torch.no_grad(): + for k in range(0, tok.shape[0], bs): + pred = model(tok[k:k + bs].to(dev)).argmax(-1).cpu() + y = lab[k:k + bs] + corr += int((pred == y).sum()); tot += int(y.shape[0]) + for c in range(len(BUCKETS)): + msk = y == c + per[c, 1] += int(msk.sum()); per[c, 0] += int(((pred == y) & msk).sum()) + return corr / max(tot, 1), per + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--deepcad", default=str(_REPO / "resources/DeepCAD/data/cad_vec")) + ap.add_argument("--n-train", type=int, default=5000) + ap.add_argument("--n-val", type=int, default=1000) + ap.add_argument("--epochs", type=int, default=30) + ap.add_argument("--bs", type=int, default=64) + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--device", default="cpu") + ap.add_argument("--out", default=str(_REPO / "ll_stepnet/checkpoints")) + args = ap.parse_args() + dev = args.device + + sys.path.insert(0, str(_REPO / "resources/DeepCAD")) + from cadlib.extrude import CADSequence + from cadlib.curves import Arc, Circle + from ll_gen.proposals.command_proposal import CommandSequenceProposal + from ll_gen.disposal.command_executor import execute_command_proposal + from stepnet.tasks import STEPForClassification + from OCC.Core.TopExp import TopExp_Explorer + from OCC.Core.TopAbs import TopAbs_FACE + + def face_count(shape): + c = 0 + e = TopExp_Explorer(shape, TopAbs_FACE) + while e.More(): + c += 1 + e.Next() + return c + + deps = (CADSequence, Circle, Arc, CommandSequenceProposal, execute_command_proposal, face_count) + files = sorted(glob.glob(os.path.join(args.deepcad, "*/*.h5"))) + if not files: + raise SystemExit(f"No cad_vec h5 files under {args.deepcad}; download DeepCAD data.tar first.") + + need = (args.n_train + args.n_val) * 3 + 4000 + print(f"building dataset from up to {need} DeepCAD models ...", flush=True) + val_tok, val_lab = load_dataset(files[:need // 6], args.n_val, deps) + tr_tok, tr_lab = load_dataset(files[need // 6:], args.n_train, deps) + print(f"built {tr_tok.shape[0]} train / {val_tok.shape[0]} val", flush=True) + print(f"train label dist: {dict(sorted(collections.Counter(tr_lab.tolist()).items()))} {CLASS_NAMES}", flush=True) + + model = STEPForClassification(num_classes=len(BUCKETS)).to(dev) + counts = np.bincount(tr_lab.numpy(), minlength=len(BUCKETS)).astype(float) + w = torch.tensor(counts.sum() / (len(BUCKETS) * np.clip(counts, 1, None)), dtype=torch.float32, device=dev) + crit = nn.CrossEntropyLoss(weight=w) + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs) + + majority = float(np.max(counts) / counts.sum()) + base_acc, _ = accuracy(model, val_tok, val_lab, dev) + print(f"BASELINE val_acc={base_acc:.3f} (majority-class baseline={majority:.3f})", flush=True) + + os.makedirs(args.out, exist_ok=True) + ckpt = os.path.join(args.out, "stepnet_classifier.pt") + n = tr_tok.shape[0] + best = -1.0 + for epoch in range(args.epochs): + model.train() + perm = torch.randperm(n) + tot = 0.0 + nb = 0 + for k in range(0, n, args.bs): + idx = perm[k:k + args.bs] + loss = crit(model(tr_tok[idx].to(dev)), tr_lab[idx].to(dev)) + opt.zero_grad() + loss.backward() + opt.step() + tot += float(loss) + nb += 1 + sched.step() + acc, _ = accuracy(model, val_tok, val_lab, dev) + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} val_acc={acc:.3f}", flush=True) + if acc > best: + best = acc + torch.save({"model_state_dict": model.state_dict(), + "num_classes": len(BUCKETS), "class_names": CLASS_NAMES}, ckpt) + + acc, per = accuracy(model, val_tok, val_lab, dev) + per_class = {CLASS_NAMES[c]: [round(per[c, 0] / max(per[c, 1], 1), 3), int(per[c, 1])] + for c in range(len(BUCKETS))} + result = {"task": "STEP->face-count complexity class (3)", "dataset": "DeepCAD cad_vec", + "n_train": int(tr_tok.shape[0]), "n_val": int(val_tok.shape[0]), "epochs": args.epochs, + "baseline_majority_acc": round(majority, 3), "best_val_acc": round(best, 3), + "final_val_acc": round(acc, 3), "per_class_acc": per_class, "checkpoint": ckpt} + with open(os.path.join(args.out, "stepnet_classifier_metrics.json"), "w") as fh: + json.dump(result, fh, indent=2) + print("DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/ll_stepnet/stepnet/diffusion.py b/ll_stepnet/stepnet/diffusion.py index 40f9819..c18888b 100644 --- a/ll_stepnet/stepnet/diffusion.py +++ b/ll_stepnet/stepnet/diffusion.py @@ -535,7 +535,9 @@ def __init__( def encode_faces(self, face_grids: torch.Tensor) -> torch.Tensor: """[B, N, U, V, 3] -> [B, N, latent_dim].""" b, n = face_grids.shape[0], face_grids.shape[1] - x = face_grids.reshape(b * n, self.uv, self.uv, 3).permute(0, 3, 1, 2) + # .contiguous() after permute: Conv2d's backward calls .view() on its + # input, which fails on the permuted (non-contiguous) tensor. + x = face_grids.reshape(b * n, self.uv, self.uv, 3).permute(0, 3, 1, 2).contiguous() z = self.face_encoder(x) return z.reshape(b, n, self.latent_dim) @@ -548,7 +550,8 @@ def decode_faces(self, latent: torch.Tensor) -> torch.Tensor: def encode_edges(self, edge_points: torch.Tensor) -> torch.Tensor: """[B, N, M, 3] -> [B, N, latent_dim].""" b, n = edge_points.shape[0], edge_points.shape[1] - x = edge_points.reshape(b * n, self.edge_pts, 3).permute(0, 2, 1) + # .contiguous() after permute (see encode_faces): Conv1d backward .view(). + x = edge_points.reshape(b * n, self.edge_pts, 3).permute(0, 2, 1).contiguous() z = self.edge_encoder(x) return z.reshape(b, n, self.latent_dim) @@ -845,8 +848,14 @@ def sample( ``face_grids`` [B, N_faces, U, V, 3] and ``edge_points`` [B, N_edges, M, 3]. """ - if device is None: - device = next(self.parameters()).device + # The denoiser weights determine where compute must happen. Always sample + # on the weights' device: a caller-supplied device that disagrees with the + # weights (e.g. a generator whose ``.device`` attribute was not updated + # after the model was ``.to()``'d) would otherwise crash at the first + # Linear with "input is on cpu but expected on mps". + param_device = next(self.parameters()).device + if device is None or torch.device(device) != param_device: + device = param_device results: Dict[str, torch.Tensor] = {} prev_denoised: Optional[torch.Tensor] = None diff --git a/ll_stepnet/stepnet/vae.py b/ll_stepnet/stepnet/vae.py index 57c6d3e..7232996 100644 --- a/ll_stepnet/stepnet/vae.py +++ b/ll_stepnet/stepnet/vae.py @@ -366,13 +366,27 @@ def forward( if param_targets is not None: param_loss = torch.tensor(0.0, device=token_ids.device) + num_contributing = 0 for i, head_logits in enumerate(param_logits): p_target = param_targets[..., i].reshape(-1) + # A parameter slot that is inactive for every command in the + # batch produces an all-``ignore_index`` target. In that + # case F.cross_entropy averages over zero elements and + # returns NaN, which poisons recon_loss and the total loss. + # The 6-command CAD schema only ever activates slots 0-7, so + # heads 8-15 are always all-ignored — without this guard the + # supervised forward is unusable. Only accumulate heads that + # carry a supervised target somewhere in the batch, and + # average over those contributing heads. + if not torch.any(p_target != -1): + continue p_logits = head_logits.reshape(-1, self.num_param_levels) param_loss = param_loss + F.cross_entropy( p_logits, p_target, ignore_index=-1 ) - recon_loss = recon_loss + param_loss / len(param_logits) + num_contributing += 1 + if num_contributing > 0: + recon_loss = recon_loss + param_loss / num_contributing outputs["recon_loss"] = recon_loss outputs["loss"] = recon_loss + self.beta * kl_loss diff --git a/ll_stepnet/tests/test_vae_sparse_param_loss.py b/ll_stepnet/tests/test_vae_sparse_param_loss.py new file mode 100644 index 0000000..b3385a0 --- /dev/null +++ b/ll_stepnet/tests/test_vae_sparse_param_loss.py @@ -0,0 +1,87 @@ +"""Regression: STEPVAE.forward must return a finite loss when some parameter +slots are inactive across the whole batch. + +The 6-command CAD schema (SOL/LINE/ARC/CIRCLE/EXTRUDE/EOS) only ever activates +parameter slots 0-7, so the param heads for slots 8-15 always receive an +all-``ignore_index`` target. Previously the supervised forward computed +``F.cross_entropy`` per head with ``reduction='mean'``; over an all-ignored +target that averages zero elements and returns NaN, poisoning ``recon_loss`` and +``loss``. This made the supervised reconstruction forward (used for DeepCAD +pretraining) unusable. The forward now skips all-ignored heads and averages +over the contributing ones. + +IMPORTANT: torch is imported by conftest.py BEFORE this module loads to avoid +OpenMP conflicts on macOS. +""" +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + + +def _build_vae(sample_encoder_config): + from stepnet import STEPVAE + + model = STEPVAE( + encoder_config=sample_encoder_config, + latent_dim=32, + max_seq_len=12, + ) + model.train() + return model + + +def test_forward_finite_loss_with_inactive_param_slots(sample_encoder_config, device): + """Slots 8-15 are never active for any command -> loss must stay finite.""" + model = _build_vae(sample_encoder_config).to(device) + + batch, seq = 2, 6 + # SOL(0), LINE(1), CIRCLE(3), EXTRUDE(4), EOS(5), pad(-1) + command_targets = torch.tensor( + [[0, 1, 3, 4, 5, -1], [0, 1, 1, 4, 5, -1]], device=device + ) + token_ids = command_targets.clamp(min=0) + attention_mask = (command_targets != -1).long() + + # Only slots 0-7 ever carry a target; slots 8-15 stay -1 everywhere. + param_targets = torch.full((batch, seq, 16), -1, dtype=torch.long, device=device) + param_targets[:, 1, 0:4] = torch.randint(0, 256, (batch, 4), device=device) # LINE + param_targets[0, 2, 0:3] = torch.randint(0, 256, (3,), device=device) # CIRCLE + param_targets[:, 3, 0:8] = torch.randint(0, 256, (batch, 8), device=device) # EXTRUDE + + out = model( + token_ids, + attention_mask=attention_mask, + command_targets=command_targets, + param_targets=param_targets, + ) + + assert torch.isfinite(out["loss"]).item(), "total loss must be finite" + assert torch.isfinite(out["recon_loss"]).item(), "recon_loss must be finite" + # The loss must be differentiable and produce real gradients. + out["loss"].backward() + grads = [p.grad for p in model.parameters() if p.grad is not None] + assert grads, "backward must populate gradients" + assert all(torch.isfinite(g).all().item() for g in grads), "gradients must be finite" + + +def test_forward_finite_when_only_one_command_type_present(sample_encoder_config, device): + """Extreme case: a batch of only CIRCLE commands (slots 3-15 all-ignored).""" + model = _build_vae(sample_encoder_config).to(device) + + command_targets = torch.tensor([[0, 3, 5], [0, 3, 5]], device=device) # SOL, CIRCLE, EOS + token_ids = command_targets.clamp(min=0) + attention_mask = torch.ones_like(command_targets) + + param_targets = torch.full((2, 3, 16), -1, dtype=torch.long, device=device) + param_targets[:, 1, 0:3] = torch.randint(0, 256, (2, 3), device=device) # CIRCLE only + + out = model( + token_ids, + attention_mask=attention_mask, + command_targets=command_targets, + param_targets=param_targets, + ) + assert torch.isfinite(out["loss"]).item() + assert torch.isfinite(out["recon_loss"]).item()