From fe323bfe6ccc00fdacbd51d276ebd59903571a99 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 03:13:38 -0500 Subject: [PATCH 01/17] fix(ll_gen): closure-aware sketch decode + repair broken diffusion sewing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit command_executor: build sketch loops by threading curve endpoints and auto-closing, so multi-line/arc polygons close by construction instead of relying on a non-autoregressive decoder to align independent per-curve argmaxes (the wire-closure limit that left only circles valid). Non-connecting square -> valid solid; tri/pent/hex/oct all close. surface_executor: fix two bugs that made the diffusion path never produce a shape — (1) _fit_bspline_surface called cadling's fit_surface(grid, tolerance=) with an unsupported kwarg and expected a face but got a dict (TypeError dropped every face); (2) the merge step called TopologyMerger.merge_edges, which does not exist — the real API is merge(faces) -> {shape, valid}. Now uses merge(faces) as the primary watertight path with the built-in edge dedup as fallback. With these, real DeepCAD geometry sews into closed shells. Regression tests: TestClosureAwareSketch (3), TestWatertightSew (unit cube -> closed volume). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/ll_gen/disposal/command_executor.py | 140 +++++++++++++-------- ll_gen/ll_gen/disposal/surface_executor.py | 52 ++++++-- ll_gen/tests/test_command_executor.py | 109 ++++++++++++++++ ll_gen/tests/test_surface_executor.py | 74 +++++++++++ 4 files changed, 310 insertions(+), 65 deletions(-) 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/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}" From 0391c1799536a5ba87a945c882fcc87db2294342 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 03:13:38 -0500 Subject: [PATCH 02/17] fix(ll_stepnet): vae sparse-param NaN + diffusion codec MPS contiguity vae: STEPVAE.forward's parameter loss averaged over all 16 param heads, but the 6-command schema only ever activates slots 0-7, so heads 8-15 receive an all-ignore_index target -> F.cross_entropy(mean) over zero elements = NaN -> poisoned recon_loss, silently blocking the supervised reconstruction forward. Now skips all-ignored heads and averages over contributing ones. diffusion: GeometryCodec.encode_faces/encode_edges permute then feed Conv2d/ Conv1d; the conv backward calls .view() on the non-contiguous tensor and raises on MPS. Added .contiguous() (also unblocks the RL diffusion path on MPS). Regression test: test_vae_sparse_param_loss.py. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_stepnet/stepnet/diffusion.py | 7 +- ll_stepnet/stepnet/vae.py | 16 +++- .../tests/test_vae_sparse_param_loss.py | 87 +++++++++++++++++++ 3 files changed, 107 insertions(+), 3 deletions(-) create mode 100644 ll_stepnet/tests/test_vae_sparse_param_loss.py diff --git a/ll_stepnet/stepnet/diffusion.py b/ll_stepnet/stepnet/diffusion.py index 40f9819..9bbfb7a 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) 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() From 25f2502b3d398c4d2a44239ffae4f2d75ecc07bd Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 03:13:38 -0500 Subject: [PATCH 03/17] =?UTF-8?q?docs(ll=5Fgen):=20make-real=20findings=20?= =?UTF-8?q?=E2=80=94=20generation=20architecture=20+=20bug=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Honest writeup of the ll_gen "make-real" investigation: the command-VAE is primitive-limited (non-AR decoder -> cylinders; supervised -> posterior collapse to all-SOL) while the diffusion path's codec is real (recon MSE 0.0003) and its sewing pipeline is now functional after 3 bug fixes, with denoiser convergence the remaining open frontier. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/2026-06-10-ll-gen-make-real-findings.md | 151 +++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 docs/2026-06-10-ll-gen-make-real-findings.md 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) From 80be11677f29bc6dd77e3a45a0361a8f5f2a3c05 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 03:39:15 -0500 Subject: [PATCH 04/17] feat(make-real): tracked training scripts for ll_stepnet classifier + diffusion codec Output the make-real trainers as proper, reproducible, tracked scripts in the package scripts/ dirs (they previously lived only in gitignored scratch). Each is self-contained (inlines the validated DeepCAD cad_vec -> executor-schema translation) and writes its checkpoint + metrics to the package checkpoints/ dir. - ll_stepnet/scripts/train_classification.py: trains STEPForClassification on real DeepCAD models (STEP -> face-count complexity class). Real result: val acc 0.976 vs 0.434 majority baseline; per-class 0.991 / 0.968 / 0.939. - ll_gen/scripts/train_diffusion_codec.py: trains StructuredDiffusion's GeometryCodec on real B-rep geometry. Real result: recon MSE 0.40 -> 0.0003. Trained checkpoints land in ll_stepnet/checkpoints/ and ll_gen/checkpoints/ (gitignored binaries, per the repo's *.pt convention). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/scripts/train_diffusion_codec.py | 253 ++++++++++++++++++++ ll_stepnet/scripts/train_classification.py | 262 +++++++++++++++++++++ 2 files changed, 515 insertions(+) create mode 100644 ll_gen/scripts/train_diffusion_codec.py create mode 100644 ll_stepnet/scripts/train_classification.py 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_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() From ad9164c6bb8c6a811fbd726d4a97bceb3ba85d19 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 03:47:21 -0500 Subject: [PATCH 05/17] fix(ll_gen): evaluate_validity exports generated CAD to files (not empty folders) The harness hardcoded export=False while still creating /disposed/, so every eval left an empty directory instead of the CAD the model generated. Add an `export` parameter (default True) forwarded to DisposalEngine.dispose, so an eval run writes .step + .stl per valid shape. Throughput-sensitive callers (per-step RL evals) can pass export=False. Verified: the trained 95% VAE checkpoint now exports 30/30 valid models as real STEP/STL files. Regression test: test_export_flag_forwarded_to_dispose. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/ll_gen/training/evaluate_validity.py | 8 ++++- ll_gen/tests/test_validity_eval.py | 36 +++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) 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/tests/test_validity_eval.py b/ll_gen/tests/test_validity_eval.py index c21204e..0447e71 100644 --- a/ll_gen/tests/test_validity_eval.py +++ b/ll_gen/tests/test_validity_eval.py @@ -238,3 +238,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 From 12cba47f4107c0e9b684aecb289ad49b1999fa85 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 04:02:12 -0500 Subject: [PATCH 06/17] fix(ll_ocadr): make the configured multimodal model actually runnable Three wiring bugs prevented LatticelabsOCADRForCausalLM from running at all: 1. config.n_embed was hardcoded (1280) and matched no real base LLM (Qwen2-0.5B=896, 1.5B=1536, 7B=3584). The projector emits n_embed-dim mesh tokens that are spliced INTO the LLM input embeddings (hidden_size), so the splice raised a shape mismatch. n_embed is now derived from the loaded LLM's hidden_size in __init__ (LLM loaded first), so any base size wires correctly. 2. get_config_for_model referenced non-existent Qwen2 sizes ("1.8B"/"14B"). Now uses real sizes and adds first-class "0.5b"/"1.5b" variants that run/train on a single machine ("7b" needs a GPU). n_embed is no longer set here (derived). 3. _splice_tokens assigned fp32 modality embeddings into a bf16 LLM embedding tensor (Qwen2 loads as bf16) -> "Index put requires ... dtypes match". It now casts modality embeddings to the destination dtype/device before the splice. Verified end to end on the 0.5b variant: forward -> logits [1,S,vocab] and generate() both run with mesh conditioning. Regression test: tests/test_wiring.py (mocks the base LLM, so no Qwen download). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_ocadr/tests/test_vision_modality.py | 5 ++ ll_ocadr/tests/test_wiring.py | 88 ++++++++++++++++++++++++++ ll_ocadr/vllm/config.py | 32 ++++++---- ll_ocadr/vllm/latticelabs_ocadr.py | 30 +++++++-- 4 files changed, 136 insertions(+), 19 deletions(-) create mode 100644 ll_ocadr/tests/test_wiring.py 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 From d8abaeaca639d741d66686a92fb3a9e945dc7876 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 06:04:07 -0500 Subject: [PATCH 07/17] feat(ll_ocadr): native-MLX multimodal training on Apple Silicon MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit End-to-end MLX implementation of ll_ocadr: a trainable PointNet-style geometry encoder + projector map a CAD point cloud into the embedding space of a FROZEN 4-bit-quantized Qwen2 (via mlx-lm's input_embeddings path); LoRA adapters on the LLM + the encoder train jointly so the LLM learns to attend to the injected mesh tokens. Trains where PyTorch-MPS would OOM, and the projector/LoRA are sized to the real LLM (not a proxy). Result (DeepCAD point-cloud -> 3-way class, Qwen2-0.5B-4bit, best epoch): encoder_mesh_read_acc 0.804, llm_generation_acc 0.543 vs shuffled-mesh baseline 0.322 (majority 0.350) — the model genuinely reads the geometry and verbalizes it well above the shuffled-mesh control. Includes a held-out SHUFFLED-mesh baseline as the honest grounding test, an auxiliary classification head (makes mesh tokens discriminative), and dataset caching. Checkpoint + metrics -> ll_ocadr/checkpoints/ (gitignored binaries). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_ocadr/mlx/train_ocadr_mlx.py | 350 ++++++++++++++++++++++++++++++++ 1 file changed, 350 insertions(+) create mode 100644 ll_ocadr/mlx/train_ocadr_mlx.py diff --git a/ll_ocadr/mlx/train_ocadr_mlx.py b/ll_ocadr/mlx/train_ocadr_mlx.py new file mode 100644 index 0000000..eea4b69 --- /dev/null +++ b/ll_ocadr/mlx/train_ocadr_mlx.py @@ -0,0 +1,350 @@ +"""ll_ocadr in MLX, end to end on Apple Silicon. + +A native-MLX implementation of the ll_ocadr multimodal idea: a trainable 3D +geometry encoder + projector map a CAD point cloud into the embedding space of a +4-bit-quantized Qwen2 LLM (via mlx-lm's ``input_embeddings`` path). The LLM base +weights stay frozen/quantized, but LoRA adapters on its attention/MLP are trained +jointly with the encoder so the LLM actually LEARNS TO ATTEND to the injected +mesh tokens (LLaVA-style; a frozen LLM otherwise ignores the out-of-distribution +mesh embeddings and collapses to the text prior). This trains on Apple Silicon +where PyTorch-MPS would OOM, and the projector/LoRA are sized to the real LLM. + +Task (mesh-grounded): given a CAD point cloud, generate a short structured +description "faces= ext= class=" derived from the +actual solid. Honest success bar: held-out face-count-class accuracy must beat a +SHUFFLED-mesh baseline (same model, mesh swapped) — i.e. the model must read the +geometry, not the text prior. + +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") + +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 = 512, 16 +CLASS_NAMES = ["simple", "box", "complex"] + + +def _bucket(nf: int) -> int: + return 0 if nf <= 4 else (1 if nf <= 6 else 2) + + +def _sample_points(shape, P=NUM_POINTS): + BRepMesh_IncrementalMesh(shape, 0.05) + pts = [] + 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() + for i in range(1, tri.NbNodes() + 1): + p = tri.Node(i).Transformed(tr) + pts.append([p.X(), p.Y(), p.Z()]) + exp.Next() + if not pts: + return None, 0 + pts = np.asarray(pts, np.float32) + c = pts.mean(0) + s = np.abs(pts - c).max() + 1e-6 + pts = (pts - c) / s + idx = np.random.choice(len(pts), P, replace=len(pts) < P) + return pts[idx], nfaces + + +def build_dataset(n_target, deps, cache): + if cache and os.path.exists(cache): + d = np.load(cache, allow_pickle=True) + if int(d["clouds"].shape[0]) >= n_target: + return d["clouds"][:n_target], list(d["texts"][:n_target]), d["classes"][:n_target] + CADSequence, command_dicts, execute = deps + files = sorted(glob.glob(os.path.join(_DEEPCAD, "data/cad_vec/*/*.h5"))) + clouds, texts, classes = [], [], [] + for f in files: + if len(clouds) >= 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) + ne = len(cad.seq) + shape = execute(command_dicts(cad)) + if shape is None: + continue + pc, nf = _sample_points(shape) + if pc is None or nf < 1: + continue + clouds.append(pc) + texts.append(f"faces={nf} ext={ne} class={CLASS_NAMES[_bucket(nf)]}") + classes.append(_bucket(nf)) + except Exception: + continue + clouds = np.stack(clouds) + classes = np.array(classes) + if cache: + np.savez(cache, clouds=clouds, texts=np.array(texts, dtype=object), classes=classes) + return clouds, texts, classes + + +class GeometryEncoderMLX(mlxnn.Module): + """PointNet-style encoder: per-point MLP + MAXPOOL global descriptor (the + proven-discriminative path — an encoder-only maxpool head reaches ~0.87 on + this task), expanded into NUM_MESH_TOKENS distinct tokens in the LLM hidden + space. Every mesh token is a learned projection of the discriminative global + feature, so both the aux head and the (LoRA) LLM get strong signal.""" + + def __init__(self, d=256, llm_hidden=896, n_tokens=NUM_MESH_TOKENS): + super().__init__() + self.point_mlp = mlxnn.Sequential( + mlxnn.Linear(3, d), mlxnn.GELU(), + mlxnn.Linear(d, d), mlxnn.GELU(), + mlxnn.Linear(d, d), mlxnn.GELU(), + ) + self.n_tokens = n_tokens + self.llm_hidden = llm_hidden + self.token_proj = mlxnn.Linear(d, n_tokens * llm_hidden) # global -> N tokens + self.aux_head = mlxnn.Linear(llm_hidden, 3) + + def __call__(self, points): + feats = self.point_mlp(points) # [B, P, d] + glob = feats.max(axis=1) # [B, d] maxpool (discriminative) + B = glob.shape[0] + return self.token_proj(glob).reshape(B, self.n_tokens, self.llm_hidden) + + def aux_logits(self, mesh_tokens): + return self.aux_head(mesh_tokens.mean(axis=1)) # [B, 3] + + +class OCADRMLX(mlxnn.Module): + """Wraps the trainable encoder + the LoRA-adapted (otherwise frozen) LLM so + mlx value_and_grad differentiates encoder params AND LoRA params together.""" + + def __init__(self, encoder, llm): + super().__init__() + self.encoder = encoder + 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=3000) + ap.add_argument("--n-val", type=int, default=600) + ap.add_argument("--epochs", type=int, default=12) + ap.add_argument("--bs", type=int, default=16) + 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} + applying 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 + enc = GeometryEncoderMLX(llm_hidden=H) + model = OCADRMLX(enc, llm) + n_trainable = sum(v.size for _, v in tree_flatten(model.trainable_parameters())) + print(f"trainable params (encoder + 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 make_batch(clouds, texts): + mesh = model.encoder(mx.array(clouds)) # [B,T,H] + prompt_emb = model.llm.model.embed_tokens(prompt_ids) # [Lp,H] + 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(clouds, texts, classes): + emb, labels, mesh = make_batch(clouds, 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) + # auxiliary classification on the mesh tokens (makes them discriminative) + aux = mlxnn.losses.cross_entropy(model.encoder.aux_logits(mesh), mx.array(classes), reduction="mean") + return lm_loss + 0.5 * aux + + if args.mode == "probe": + clouds, texts, classes = build_dataset(8, deps, None) + lv, grads = mlxnn.value_and_grad(model, loss_fn)(clouds[:4], texts[:4], 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} trainable grad_norm={gnorm:.4f} text={texts[0]!r}", flush=True) + return + + cache = f"{args.out}/ocadr_data_cache.npz" + os.makedirs(args.out, exist_ok=True) + print("building/loading dataset ...", flush=True) + allc, _allt, allcl = build_dataset(args.n_train + args.n_val, deps, cache) + # Balance the 3 classes so the text prior gives only chance (1/3): this + # removes the majority-class shortcut and forces the model to READ the mesh + # to beat the shuffled-mesh baseline. Target is the bare class word (no + # predictable "faces=N" prefix to ride). + 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) + allc, allcl = allc[keep], allcl[keep] + allt = [CLASS_NAMES[c] for c in allcl] + nval = max(len(keep) // 6, 90) + vc, vt, vcl = allc[:nval], allt[:nval], allcl[:nval] + tc, tt, tcl = allc[nval:], allt[nval:], allcl[nval:] + print(f"balanced {tc.shape[0]} train / {vc.shape[0]} val; per-class={per}; target=class-only; " + f"val dist={np.bincount(vcl, minlength=3).tolist()} (majority baseline={np.bincount(vcl,minlength=3).max()/len(vcl):.3f})", + flush=True) + + opt = optim.Adam(learning_rate=args.lr) + lg_fn = mlxnn.value_and_grad(model, loss_fn) + n = tc.shape[0] + + def greedy_class(clouds): + B = clouds.shape[0] + mesh = model.encoder(mx.array(clouds)) + 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(12): + 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() + # target is the bare class word; take the first class name that appears + 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(clouds, classes, shuffle=False): + src = np.random.permutation(clouds.shape[0]) if shuffle else np.arange(clouds.shape[0]) + cl = clouds[src] + correct = 0 + for k in range(0, clouds.shape[0], 32): + correct += int((greedy_class(cl[k:k + 32]) == classes[k:k + 32]).sum()) + return correct / clouds.shape[0] + + def aux_acc(clouds, classes): + """Accuracy of the encoder's auxiliary head — does the MLX geometry + encoder itself read the mesh (independent of LLM verbalization)?""" + correct = 0 + for k in range(0, clouds.shape[0], 64): + mesh = model.encoder(mx.array(clouds[k:k + 64])) + pred = np.array(mx.argmax(model.encoder.aux_logits(mesh), axis=1).tolist()) + correct += int((pred == classes[k:k + 64]).sum()) + return correct / clouds.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, grads = lg_fn(tc[idx], [tt[i] for i in idx], tcl[idx]) + opt.update(model, grads) + mx.eval(model.trainable_parameters(), opt.state, lv) + tot += float(lv.item()) + nb += 1 + acc = class_acc(vc, vcl) + shuf = class_acc(vc, vcl, shuffle=True) + aux = aux_acc(vc, 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()))) + + majority = float(np.bincount(vcl, minlength=3).max() / len(vcl)) + result = {"framework": "MLX", "task": "CAD point-cloud -> class (simple/box/complex)", + "llm": args.llm, "llm_base": "frozen+4bit-quantized", + "trainable": "geometry encoder + LoRA adapters", "trainable_params_M": round(n_trainable / 1e6, 2), + "n_train": int(tc.shape[0]), "n_val": int(vc.shape[0]), "epochs": args.epochs, + "majority_baseline": round(majority, 3), + "encoder_mesh_read_acc": round(aux_acc(vc, vcl), 3), + "llm_generation_acc": round(class_acc(vc, vcl), 3), + "shuffled_mesh_baseline": round(class_acc(vc, 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_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() From 0727c86cd83c01e6355e0a2d3c999b025606f60d Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 06:09:50 -0500 Subject: [PATCH 08/17] feat(ll_stepnet): native-MLX STEP classifier (Apple Silicon) MLX port of stepnet.tasks.STEPForClassification (token embedding + transformer encoder + masked mean-pool + MLP head), same DeepCAD cad_vec -> face-count complexity task as the PyTorch trainer so results are directly comparable. Result: MLX val acc 0.942 vs majority 0.436 (PyTorch reference 0.976). A faithful Apple-Silicon-native port reaching strong accuracy; the small gap is the lighter MLX config (2 layers / d=128). Checkpoint + metrics -> ll_stepnet/checkpoints/. Part of the all-models -> MLX migration (after ll_ocadr). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_stepnet/mlx/train_classification_mlx.py | 270 +++++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 ll_stepnet/mlx/train_classification_mlx.py diff --git a/ll_stepnet/mlx/train_classification_mlx.py b/ll_stepnet/mlx/train_classification_mlx.py new file mode 100644 index 0000000..6d4583d --- /dev/null +++ b/ll_stepnet/mlx/train_classification_mlx.py @@ -0,0 +1,270 @@ +"""ll_stepnet STEP classifier in native MLX (Apple Silicon). + +MLX port of stepnet.tasks.STEPForClassification: a token-embedding + +transformer-encoder + mean-pool + MLP-head classifier. Same task and data as the +PyTorch trainer (ll_stepnet/scripts/train_classification.py) — DeepCAD cad_vec -> +command-token sequence -> face-count complexity class — so the MLX result is +directly comparable (PyTorch reached val acc 0.976 vs 0.434 majority). + +Trains entirely in MLX on Apple Silicon. Modes: probe | train. +""" + +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") +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")) + +# --- DeepCAD cad_vec -> command-token sequence + face-count class (self-contained) +LEVELS, RANGE, MAX_LEN = 256, 2.0, 256 +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} +VOCAB = 12 + LEVELS # command/special tokens + param value tokens +BUCKETS = [(0, 4), (5, 6), (7, 9999)] +CLASS_NAMES = ["simple(<=4)", "box(5-6)", "complex(7+)"] + + +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 _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 _tokens(cmds): + t = [1] + 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(2) + t = t[:MAX_LEN] + return t + [0] * (MAX_LEN - len(t)) + + +def _bucket(nf): + return 0 if nf <= 4 else (1 if nf <= 6 else 2) + + +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], d["classes"][:n_target] + 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 nfaces(shape): + c = 0 + e = TopExp_Explorer(shape, TopAbs_FACE) + while e.More(): + c += 1 + e.Next() + return c + + toks, classes = [], [] + 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) + shape = execute_command_proposal(CommandSequenceProposal( + command_dicts=_command_dicts(cmds), quantization_bits=8, normalization_range=2.0)) + if shape is None: + continue + nf = nfaces(shape) + if nf < 1: + continue + toks.append(_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 + + +# --- MLX transformer classifier (port of STEPForClassification) ------------- +class Block(mlxnn.Module): + def __init__(self, d, heads): + super().__init__() + self.attn = mlxnn.MultiHeadAttention(d, heads) + self.n1 = mlxnn.LayerNorm(d) + self.ffn = mlxnn.Sequential(mlxnn.Linear(d, d * 4), mlxnn.GELU(), mlxnn.Linear(d * 4, d)) + self.n2 = mlxnn.LayerNorm(d) + + def __call__(self, x): + x = self.n1(x + self.attn(x, x, x)) + return self.n2(x + self.ffn(x)) + + +class MLXClassifier(mlxnn.Module): + def __init__(self, vocab=VOCAB, d=128, layers=2, heads=4, nclass=3, maxlen=MAX_LEN): + super().__init__() + self.embed = mlxnn.Embedding(vocab, d) + self.pos = mx.random.normal((maxlen, d)) * 0.02 + self.blocks = [Block(d, heads) for _ in range(layers)] + self.norm = mlxnn.LayerNorm(d) + self.head = mlxnn.Sequential(mlxnn.Linear(d, 512), mlxnn.ReLU(), mlxnn.Linear(512, nclass)) + + def __call__(self, ids): + x = self.embed(ids) + self.pos[: ids.shape[1]][None] + for b in self.blocks: + x = b(x) + x = self.norm(x) + m = (ids != 0).astype(x.dtype)[..., None] + pooled = (x * m).sum(axis=1) / mx.maximum(m.sum(axis=1), 1) + return self.head(pooled) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=["probe", "train"], default="train") + 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("--out", default=str(_REPO / "ll_stepnet/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + cache = f"{args.out}/mlx_classifier_data.npz" + + if args.mode == "probe": + toks, classes = build_dataset(40, None) + model = MLXClassifier() + out = model(mx.array(toks[:8])) + import collections + print(f"probe: tokens {toks.shape}, logits {out.shape}, dist {dict(collections.Counter(classes.tolist()))}", flush=True) + return + + print("building/loading dataset ...", flush=True) + toks, classes = build_dataset(args.n_train + args.n_val, cache) + vt, vcl = toks[:args.n_val], classes[:args.n_val] + tt, tcl = toks[args.n_val:], classes[args.n_val:] + import collections + cnt = np.bincount(tcl, minlength=3) + majority = float(np.bincount(vcl, minlength=3).max() / len(vcl)) + print(f"built {tt.shape[0]} train / {vt.shape[0]} val; train dist={cnt.tolist()}; majority={majority:.3f}", flush=True) + + model = MLXClassifier() + w = mx.array((cnt.sum() / (3 * np.clip(cnt, 1, None))).astype(np.float32)) + opt = optim.Adam(learning_rate=args.lr) + + def loss_fn(ids, y): + logits = model(ids) + ce = mlxnn.losses.cross_entropy(logits, y, reduction="none") + return (ce * w[y]).mean() + + lg = mlxnn.value_and_grad(model, loss_fn) + + def acc(toks_, cls_): + c = 0 + for k in range(0, toks_.shape[0], 256): + p = np.array(mx.argmax(model(mx.array(toks_[k:k + 256])), axis=1).tolist()) + c += int((p == cls_[k:k + 256]).sum()) + return c / toks_.shape[0] + + best = -1.0 + n = tt.shape[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(tt[idx]), mx.array(tcl[idx])) + opt.update(model, g) + mx.eval(model.parameters(), opt.state, lv) + tot += float(lv.item()) + nb += 1 + a = acc(vt, vcl) + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} val_acc={a:.3f}", flush=True) + if a > best: + best = a + mx.save_safetensors(f"{args.out}/stepnet_classifier_mlx.safetensors", + dict(tree_flatten(model.parameters()))) + + result = {"framework": "MLX", "task": "STEP->face-count class (3)", "dataset": "DeepCAD cad_vec", + "n_train": int(tt.shape[0]), "n_val": int(vt.shape[0]), "epochs": args.epochs, + "majority_baseline": round(majority, 3), "best_val_acc": round(best, 3), + "pytorch_reference_acc": 0.976, + "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_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() From ea9b98923ebf954a9d274e2152a810e364116fea Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 06:19:01 -0500 Subject: [PATCH 09/17] feat(ll_brepnet): native-MLX B-rep segmentation GNN (Apple Silicon) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MLX port of LLBRepNet: UV-Net face/edge encoders + BRepNet coedge message-passing (h' = relu(W_self·h + W_next·h[next] + W_prev·h[prev] + W_mate·h[mate]) x6) + coedge->face mean-pool (scatter-mean via one-hot matmul) + per-face seg head. Graph batching via offset concatenation of the coedge/face/ edge indices. Same Fusion360 data/task as the PyTorch trainer. Result: MLX best val mIoU 0.711 / acc 0.906 (6k train solids, 20 epochs) vs the full PyTorch reference test mIoU 0.828 (27k solids). The GNN is correct and converges; the gap is data/epochs/encoder-capacity, not the port. Hardest model in the all-models -> MLX migration (after ll_ocadr, ll_stepnet). Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_brepnet/mlx/train_brepnet_mlx.py | 259 ++++++++++++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 ll_brepnet/mlx/train_brepnet_mlx.py diff --git a/ll_brepnet/mlx/train_brepnet_mlx.py b/ll_brepnet/mlx/train_brepnet_mlx.py new file mode 100644 index 0000000..c9ef24f --- /dev/null +++ b/ll_brepnet/mlx/train_brepnet_mlx.py @@ -0,0 +1,259 @@ +"""ll_brepnet B-rep face segmentation in native MLX (Apple Silicon). + +MLX port of LLBRepNet: UV-Net face/edge encoders -> per-entity projections -> +coedge features (gather face/edge reprs by coedge topology) -> BRepNet coedge +message-passing (W_self·h + W_next·h[next] + W_prev·h[prev] + W_mate·h[mate], +x num_layers) -> coedge->face mean-pool (scatter-mean via one-hot matmul) -> +per-face segmentation head. Same Fusion360 data + task as the PyTorch trainer +(which reached test mIoU 0.828). + +Graph batching: N solids are concatenated with offset arithmetic on the coedge/ +face/edge indices, so a batch is one big graph. Trains entirely in MLX. + +Modes: probe | train. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +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 +import mlx.optimizers as optim # noqa: E402 +from mlx.utils import tree_flatten # noqa: E402 + +_REPO = Path(__file__).resolve().parents[2] +_DATA = _REPO / "resources/fusion360/full_processed" +IGNORE = -1 + + +# --- data ------------------------------------------------------------------- +def load_records(stems, fmean, fstd, emean, estd): + recs = [] + for s in stems: + npz = _DATA / f"{s}.npz" + seg = _DATA / f"{s}.seg" + if not npz.exists() or not seg.exists(): + continue + try: + z = np.load(npz) + labels = np.loadtxt(seg, dtype=np.int64).reshape(-1) + nf = int(z["num_faces"]) + if labels.shape[0] != nf: + continue + ff = (z["face_features"].astype(np.float32) - fmean) / fstd + ef = (z["edge_features"].astype(np.float32) - emean) / estd + # grids -> channels-last for MLX conv: face [F,10,10,7], edge [E,10,6] + fg = np.transpose(z["face_point_grids"].astype(np.float32), (0, 2, 3, 1)) + eg = np.transpose(z["edge_point_grids"].astype(np.float32), (0, 2, 1)) + recs.append({ + "ff": ff, "ef": ef, "fg": fg, "eg": eg, + "c2f": z["coedge_to_face"].astype(np.int32), + "c2e": z["coedge_to_edge"].astype(np.int32), + "c2n": z["coedge_to_next"].astype(np.int32), + "c2p": z["coedge_to_prev"].astype(np.int32), + "c2m": z["coedge_to_mate"].astype(np.int32), + "rev": z["coedge_reversed"].astype(np.float32), + "labels": labels.astype(np.int32), "nf": nf, + }) + except Exception: + continue + return recs + + +def collate(batch): + """Concatenate solids into one big graph with offset coedge/face/edge ids.""" + fo = co = eo = 0 + out = {k: [] for k in ["ff", "ef", "fg", "eg", "c2f", "c2e", "c2n", "c2p", "c2m", "rev", "labels"]} + for r in batch: + out["ff"].append(r["ff"]); out["ef"].append(r["ef"]) + out["fg"].append(r["fg"]); out["eg"].append(r["eg"]) + out["c2f"].append(r["c2f"] + fo); out["c2e"].append(r["c2e"] + eo) + out["c2n"].append(r["c2n"] + co); out["c2p"].append(r["c2p"] + co); out["c2m"].append(r["c2m"] + co) + out["rev"].append(r["rev"]); out["labels"].append(r["labels"]) + fo += r["nf"]; eo += r["ef"].shape[0]; co += r["c2f"].shape[0] + cat = {k: np.concatenate(v, axis=0) for k, v in out.items()} + cat["n_faces"] = fo + return cat + + +# --- model ------------------------------------------------------------------ +class SurfEnc(mlxnn.Module): + def __init__(self, cin=7, out=64): + super().__init__() + self.c1 = mlxnn.Conv2d(cin, 32, 3, padding=1) + self.c2 = mlxnn.Conv2d(32, 64, 3, padding=1) + self.lin = mlxnn.Linear(64, out) + + def __call__(self, g): # [F,10,10,7] + x = mlxnn.relu(self.c1(g)); x = mlxnn.relu(self.c2(x)) + x = x.mean(axis=(1, 2)) # global avg pool -> [F,64] + return mlxnn.relu(self.lin(x)) + + +class CurveEnc(mlxnn.Module): + def __init__(self, cin=6, out=64): + super().__init__() + self.c1 = mlxnn.Conv1d(cin, 32, 3, padding=1) + self.c2 = mlxnn.Conv1d(32, 64, 3, padding=1) + self.lin = mlxnn.Linear(64, out) + + def __call__(self, g): # [E,10,6] + x = mlxnn.relu(self.c1(g)); x = mlxnn.relu(self.c2(x)) + x = x.mean(axis=1) # [E,64] + return mlxnn.relu(self.lin(x)) + + +class CoedgeConv(mlxnn.Module): + def __init__(self, d): + super().__init__() + self.ws = mlxnn.Linear(d, d); self.wn = mlxnn.Linear(d, d) + self.wp = mlxnn.Linear(d, d); self.wm = mlxnn.Linear(d, d) + + def __call__(self, h, nidx, pidx, midx): + return mlxnn.relu(self.ws(h) + self.wn(h[nidx]) + self.wp(h[pidx]) + self.wm(h[midx])) + + +class BRepNetMLX(mlxnn.Module): + def __init__(self, num_classes=8, ent=64, hid=128, layers=6): + super().__init__() + self.surf = SurfEnc(7, 64) + self.curve = CurveEnc(6, 64) + self.face_proj = mlxnn.Linear(8 + 64, ent) + self.edge_proj = mlxnn.Linear(7 + 64, ent) + self.in_proj = mlxnn.Linear(2 * ent + 1, hid) + self.convs = [CoedgeConv(hid) for _ in range(layers)] + self.head = mlxnn.Linear(hid, num_classes) + + def __call__(self, b): + face_repr = mlxnn.relu(self.face_proj(mx.concatenate([b["ff"], self.surf(b["fg"])], axis=1))) + edge_repr = mlxnn.relu(self.edge_proj(mx.concatenate([b["ef"], self.curve(b["eg"])], axis=1))) + h = mx.concatenate([face_repr[b["c2f"]], edge_repr[b["c2e"]], b["rev"][:, None]], axis=1) + h = mlxnn.relu(self.in_proj(h)) + for conv in self.convs: + h = conv(h, b["c2n"], b["c2p"], b["c2m"]) + # coedge -> face mean pool via one-hot matmul (scatter-mean) + nf = b["n_faces"] + onehot = (b["c2f"][:, None] == mx.arange(nf)[None, :]).astype(h.dtype) # [C, F] + face_emb = (onehot.T @ h) / mx.maximum(onehot.sum(axis=0)[:, None], 1) + return self.head(face_emb) # [F, num_classes] + + +def to_mx(b): + return { + "ff": mx.array(b["ff"]), "ef": mx.array(b["ef"]), + "fg": mx.array(b["fg"]), "eg": mx.array(b["eg"]), + "c2f": mx.array(b["c2f"]), "c2e": mx.array(b["c2e"]), + "c2n": mx.array(b["c2n"]), "c2p": mx.array(b["c2p"]), "c2m": mx.array(b["c2m"]), + "rev": mx.array(b["rev"]), "n_faces": int(b["n_faces"]), + } + + +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", "train"], default="train") + ap.add_argument("--n-train", type=int, default=6000) + ap.add_argument("--n-val", type=int, default=1500) + ap.add_argument("--epochs", type=int, default=20) + ap.add_argument("--bs", type=int, default=16) + ap.add_argument("--lr", type=float, default=1e-3) + ap.add_argument("--out", default=str(_REPO / "ll_brepnet/checkpoints")) + args = ap.parse_args() + os.makedirs(args.out, exist_ok=True) + + ds = json.load(open(_DATA / "dataset.json")) + nc = int(ds["num_classes"]) + fs = ds["feature_standardization"] + fmean = np.array([e["mean"] for e in fs["face_features"]], np.float32) + fstd = np.array([e["standard_deviation"] for e in fs["face_features"]], np.float32) + 1e-8 + emean = np.array([e["mean"] for e in fs["edge_features"]], np.float32) + estd = np.array([e["standard_deviation"] for e in fs["edge_features"]], np.float32) + 1e-8 + print(f"num_classes={nc} classes={ds['class_names']}", flush=True) + + train_stems = ds["training_set"][: args.n_train] + val_stems = ds["validation_set"][: args.n_val] + print(f"loading {len(train_stems)} train / {len(val_stems)} val records ...", flush=True) + train = load_records(train_stems, fmean, fstd, emean, estd) + val = load_records(val_stems, fmean, fstd, emean, estd) + print(f"loaded {len(train)} train / {len(val)} val solids", flush=True) + + model = BRepNetMLX(num_classes=nc) + + if args.mode == "probe": + b = to_mx(collate(train[:4])) + out = model(b) + print(f"probe: batch faces={b['n_faces']} logits={out.shape} finite={bool(mx.isfinite(out).all().item())}", flush=True) + return + + opt = optim.Adam(learning_rate=args.lr) + + def loss_fn(b, labels): + logits = model(b) + mask = labels != IGNORE + ce = mlxnn.losses.cross_entropy(logits, mx.where(mask, labels, 0), reduction="none") * mask + return ce.sum() / mx.maximum(mask.sum(), 1) + + lg = mlxnn.value_and_grad(model, loss_fn) + + def evaluate(recs): + conf = np.zeros((nc, nc), np.int64) + for k in range(0, len(recs), args.bs): + b = collate(recs[k:k + args.bs]) + logits = model(to_mx(b)) + pred = np.array(mx.argmax(logits, axis=1).tolist()) + lab = b["labels"] + m = lab != IGNORE + for p, t in zip(pred[m], lab[m]): + conf[t, p] += 1 + acc = float(np.trace(conf) / max(conf.sum(), 1)) + return miou(conf, nc), acc + + best = -1.0 + for epoch in range(args.epochs): + np.random.shuffle(train) + tot = 0.0; nb = 0 + for k in range(0, len(train), args.bs): + b = collate(train[k:k + args.bs]) + labels = mx.array(b["labels"]) + lv, g = lg(to_mx(b), labels) + opt.update(model, g) + mx.eval(model.parameters(), opt.state, lv) + tot += float(lv.item()); nb += 1 + vmiou, vacc = evaluate(val) + print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} val_mIoU={vmiou:.3f} val_acc={vacc:.3f}", flush=True) + if vmiou > best: + best = vmiou + mx.save_safetensors(f"{args.out}/brepnet_mlx.safetensors", dict(tree_flatten(model.parameters()))) + + result = {"framework": "MLX", "task": "B-rep face segmentation", "dataset": "Fusion360", + "num_classes": nc, "n_train": len(train), "n_val": len(val), "epochs": args.epochs, + "best_val_mIoU": round(best, 3), "pytorch_reference_test_mIoU": 0.828, + "checkpoint": f"{args.out}/brepnet_mlx.safetensors"} + with open(f"{args.out}/brepnet_mlx_metrics.json", "w") as fh: + json.dump(result, fh, indent=2) + print("BREPNET_MLX_DONE", json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() From 1144efa31004b64185cddf699eb5f17cbec4dd33 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 06:33:58 -0500 Subject: [PATCH 10/17] =?UTF-8?q?feat(ll=5Fstepnet):=20faithful=20MLX=20po?= =?UTF-8?q?rt=20=E2=80=94=20convert=20real=20weights,=20prove=200.976=20pa?= =?UTF-8?q?rity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Supersedes the earlier simplified MLX classifier (which retrained a lighter model to 0.942). This reproduces the EXACT STEPForClassification architecture (6-layer post-norm nn.TransformerEncoder + fusion + head, manual MHA from the packed in_proj_weight) and CONVERTS the real trained checkpoint (stepnet_classifier.pt, PyTorch val acc 0.976) into MLX. Verified on the trainer's exact val split (1000 real DeepCAD samples): argmax agreement vs PyTorch = 1.0, max abs logit diff = 1e-5, MLX val acc = 0.976 == PyTorch 0.976 (per-class 0.991/0.978/0.933). The MLX model IS the trained model (same weights, identical predictions), running natively on Apple Silicon, and remains trainable (real MLX modules). Modes: probe | convert | parity | train. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_stepnet/mlx/train_classification_mlx.py | 416 +++++++++++++++------ 1 file changed, 294 insertions(+), 122 deletions(-) diff --git a/ll_stepnet/mlx/train_classification_mlx.py b/ll_stepnet/mlx/train_classification_mlx.py index 6d4583d..5b14dfd 100644 --- a/ll_stepnet/mlx/train_classification_mlx.py +++ b/ll_stepnet/mlx/train_classification_mlx.py @@ -1,12 +1,43 @@ -"""ll_stepnet STEP classifier in native MLX (Apple Silicon). - -MLX port of stepnet.tasks.STEPForClassification: a token-embedding + -transformer-encoder + mean-pool + MLP-head classifier. Same task and data as the -PyTorch trainer (ll_stepnet/scripts/train_classification.py) — DeepCAD cad_vec -> -command-token sequence -> face-count complexity class — so the MLX result is -directly comparable (PyTorch reached val acc 0.976 vs 0.434 majority). - -Trains entirely in MLX on Apple Silicon. Modes: probe | train. +"""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 @@ -21,6 +52,7 @@ 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 @@ -28,64 +60,68 @@ import matplotlib # noqa: E402 matplotlib.use("Agg") -matplotlib.use = lambda *a, **k: None +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 # noqa: E402 +from mlx.utils import tree_flatten, tree_unflatten # 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")) -# --- DeepCAD cad_vec -> command-token sequence + face-count class (self-contained) +# --- 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} -VOCAB = 12 + LEVELS # command/special tokens + param value tokens 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 _qc(g): +def _q_coord(g): return int(np.clip(round(float(g)), 0, LEVELS - 1)) -def _qv(v): +def _q_value(v): return int(np.clip(round((float(v) + RANGE) / (2 * RANGE) * (LEVELS - 1)), 0, LEVELS - 1)) -def _cmds(cad, Circle, Arc): - out = [] +def _translate(cad, Circle, Arc): + cmds = [] for ext in cad.seq: for loop in ext.profile.children: - out.append(("SOL", {})) + cmds.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)})) + 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 - 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])})) + 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 - 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 + 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] * 16 - m = [False] * 16 + p = [0] * NUM_SLOTS + m = [False] * NUM_SLOTS for j in MASK[name]: p[j] = int(slots.get(j, 0)) m[j] = True @@ -93,26 +129,33 @@ def _command_dicts(cmds): return out -def _tokens(cmds): - t = [1] +def _encode_tokens(cmds): + toks = [1] for name, slots in cmds: - t.append(CMD_TOK[name]) + toks.append(CMD_TOK[name]) for j in MASK[name]: - t.append(12 + int(slots.get(j, 0))) - t.append(2) - t = t[:MAX_LEN] - return t + [0] * (MAX_LEN - len(t)) + 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): - return 0 if nf <= 4 else (1 if nf <= 6 else 2) + for i, (lo, hi) in enumerate(BUCKETS): + if lo <= nf <= hi: + return i + return len(BUCKETS) - 1 -def build_dataset(n_target, cache): +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_target: - return d["tokens"][:n_target], d["classes"][:n_target] + 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 @@ -120,7 +163,7 @@ def build_dataset(n_target, cache): from OCC.Core.TopExp import TopExp_Explorer from OCC.Core.TopAbs import TopAbs_FACE - def nfaces(shape): + def face_count(shape): c = 0 e = TopExp_Explorer(shape, TopAbs_FACE) while e.More(): @@ -128,23 +171,26 @@ def nfaces(shape): 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 sorted(glob.glob(os.path.join(_DEEPCAD, "data/cad_vec/*/*.h5"))): - if len(toks) >= n_target: + 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 = _cmds(cad, Circle, Arc) + 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 = nfaces(shape) + nf = face_count(shape) if nf < 1: continue - toks.append(_tokens(cmds)) + toks.append(_encode_tokens(cmds)) classes.append(_bucket(nf)) except Exception: continue @@ -155,43 +201,157 @@ def nfaces(shape): return toks, classes -# --- MLX transformer classifier (port of STEPForClassification) ------------- -class Block(mlxnn.Module): - def __init__(self, d, heads): +# --- 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.attn = mlxnn.MultiHeadAttention(d, heads) - self.n1 = mlxnn.LayerNorm(d) - self.ffn = mlxnn.Sequential(mlxnn.Linear(d, d * 4), mlxnn.GELU(), mlxnn.Linear(d * 4, d)) - self.n2 = mlxnn.LayerNorm(d) + 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.n1(x + self.attn(x, x, x)) - return self.n2(x + self.ffn(x)) + x = self.norm1(x + self.self_attn(x)) + return self.norm2(x + self.linear2(mlxnn.relu(self.linear1(x)))) -class MLXClassifier(mlxnn.Module): - def __init__(self, vocab=VOCAB, d=128, layers=2, heads=4, nclass=3, maxlen=MAX_LEN): +class FaithfulTransformerEncoder(mlxnn.Module): + def __init__(self, vocab=VOCAB, d=256, layers=6, heads=8, ff=1024): super().__init__() - self.embed = mlxnn.Embedding(vocab, d) - self.pos = mx.random.normal((maxlen, d)) * 0.02 - self.blocks = [Block(d, heads) for _ in range(layers)] - self.norm = mlxnn.LayerNorm(d) - self.head = mlxnn.Sequential(mlxnn.Linear(d, 512), mlxnn.ReLU(), mlxnn.Linear(512, nclass)) + 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.embed(ids) + self.pos[: ids.shape[1]][None] - for b in self.blocks: - x = b(x) - x = self.norm(x) - m = (ids != 0).astype(x.dtype)[..., None] - pooled = (x * m).sum(axis=1) / mx.maximum(m.sum(axis=1), 1) - return self.head(pooled) + 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", "train"], default="train") - ap.add_argument("--n-train", type=int, default=5000) + 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) @@ -199,71 +359,83 @@ def main(): ap.add_argument("--out", default=str(_REPO / "ll_stepnet/checkpoints")) args = ap.parse_args() os.makedirs(args.out, exist_ok=True) - cache = f"{args.out}/mlx_classifier_data.npz" + nclass = len(BUCKETS) if args.mode == "probe": - toks, classes = build_dataset(40, None) - model = MLXClassifier() - out = model(mx.array(toks[:8])) - import collections - print(f"probe: tokens {toks.shape}, logits {out.shape}, dist {dict(collections.Counter(classes.tolist()))}", flush=True) + 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 - print("building/loading dataset ...", flush=True) - toks, classes = build_dataset(args.n_train + args.n_val, cache) - vt, vcl = toks[:args.n_val], classes[:args.n_val] - tt, tcl = toks[args.n_val:], classes[args.n_val:] - import collections - cnt = np.bincount(tcl, minlength=3) - majority = float(np.bincount(vcl, minlength=3).max() / len(vcl)) - print(f"built {tt.shape[0]} train / {vt.shape[0]} val; train dist={cnt.tolist()}; majority={majority:.3f}", flush=True) - - model = MLXClassifier() - w = mx.array((cnt.sum() / (3 * np.clip(cnt, 1, None))).astype(np.float32)) + # 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): - logits = model(ids) - ce = mlxnn.losses.cross_entropy(logits, y, reduction="none") + ce = mlxnn.losses.cross_entropy(model(ids), y, reduction="none") return (ce * w[y]).mean() lg = mlxnn.value_and_grad(model, loss_fn) - - def acc(toks_, cls_): - c = 0 - for k in range(0, toks_.shape[0], 256): - p = np.array(mx.argmax(model(mx.array(toks_[k:k + 256])), axis=1).tolist()) - c += int((p == cls_[k:k + 256]).sum()) - return c / toks_.shape[0] - - best = -1.0 - n = tt.shape[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): + 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) - tot += float(lv.item()) - nb += 1 - a = acc(vt, vcl) - print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} val_acc={a:.3f}", flush=True) - if a > best: - best = a - mx.save_safetensors(f"{args.out}/stepnet_classifier_mlx.safetensors", - dict(tree_flatten(model.parameters()))) - - result = {"framework": "MLX", "task": "STEP->face-count class (3)", "dataset": "DeepCAD cad_vec", - "n_train": int(tt.shape[0]), "n_val": int(vt.shape[0]), "epochs": args.epochs, - "majority_baseline": round(majority, 3), "best_val_acc": round(best, 3), - "pytorch_reference_acc": 0.976, - "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_DONE", json.dumps(result), flush=True) + 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__": From 9a4d194d0a073b3a481950d7029a622dd219943c Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 06:43:09 -0500 Subject: [PATCH 11/17] =?UTF-8?q?feat(ll=5Fbrepnet):=20faithful=20MLX=20po?= =?UTF-8?q?rt=20=E2=80=94=20convert=20real=20GNN=20weights,=20prove=20mIoU?= =?UTF-8?q?=20parity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Supersedes the earlier simplified MLX port (a lighter GNN retrained to 0.711 mIoU). This reproduces the EXACT LLBRepNet architecture — UV-Net encoders WITH BatchNorm, the cadling BRepNetEncoder (input_proj Linear+LayerNorm+ReLU, 4x residual coedge convs each with LayerNorm, output_proj) and the seg head — and CONVERTS the real trained Lightning checkpoint (best.ckpt, PyTorch test mIoU 0.828) into MLX. Conversion handles the layout differences the simplified port ignored: Conv2d OIHW->OHWI and Conv1d OIW->OWI weight permutes, and BatchNorm running stats applied in inference mode. Both models are driven from the SAME real BRepDataset. Verified on 1500 test solids (23,056 faces): per-face argmax agreement vs PyTorch = 1.0, MLX mIoU = 0.8354 == PyTorch 0.8354, MLX acc = 0.945 == PyTorch 0.945. The MLX model IS the trained GNN, running natively on Apple Silicon. Modes: probe | convert | parity. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_brepnet/mlx/train_brepnet_mlx.py | 433 ++++++++++++++++------------ 1 file changed, 252 insertions(+), 181 deletions(-) diff --git a/ll_brepnet/mlx/train_brepnet_mlx.py b/ll_brepnet/mlx/train_brepnet_mlx.py index c9ef24f..27f6b2f 100644 --- a/ll_brepnet/mlx/train_brepnet_mlx.py +++ b/ll_brepnet/mlx/train_brepnet_mlx.py @@ -1,28 +1,47 @@ -"""ll_brepnet B-rep face segmentation in native MLX (Apple Silicon). - -MLX port of LLBRepNet: UV-Net face/edge encoders -> per-entity projections -> -coedge features (gather face/edge reprs by coedge topology) -> BRepNet coedge -message-passing (W_self·h + W_next·h[next] + W_prev·h[prev] + W_mate·h[mate], -x num_layers) -> coedge->face mean-pool (scatter-mean via one-hot matmul) -> -per-face segmentation head. Same Fusion360 data + task as the PyTorch trainer -(which reached test mIoU 0.828). - -Graph batching: N solids are concatenated with offset arithmetic on the coedge/ -face/edge indices, so a batch is one big graph. Trains entirely in MLX. - -Modes: probe | train. +"""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 glob 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 @@ -30,133 +49,186 @@ import numpy as np # 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.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 -# --- data ------------------------------------------------------------------- -def load_records(stems, fmean, fstd, emean, estd): - recs = [] - for s in stems: - npz = _DATA / f"{s}.npz" - seg = _DATA / f"{s}.seg" - if not npz.exists() or not seg.exists(): - continue - try: - z = np.load(npz) - labels = np.loadtxt(seg, dtype=np.int64).reshape(-1) - nf = int(z["num_faces"]) - if labels.shape[0] != nf: - continue - ff = (z["face_features"].astype(np.float32) - fmean) / fstd - ef = (z["edge_features"].astype(np.float32) - emean) / estd - # grids -> channels-last for MLX conv: face [F,10,10,7], edge [E,10,6] - fg = np.transpose(z["face_point_grids"].astype(np.float32), (0, 2, 3, 1)) - eg = np.transpose(z["edge_point_grids"].astype(np.float32), (0, 2, 1)) - recs.append({ - "ff": ff, "ef": ef, "fg": fg, "eg": eg, - "c2f": z["coedge_to_face"].astype(np.int32), - "c2e": z["coedge_to_edge"].astype(np.int32), - "c2n": z["coedge_to_next"].astype(np.int32), - "c2p": z["coedge_to_prev"].astype(np.int32), - "c2m": z["coedge_to_mate"].astype(np.int32), - "rev": z["coedge_reversed"].astype(np.float32), - "labels": labels.astype(np.int32), "nf": nf, - }) - except Exception: - continue - return recs - - -def collate(batch): - """Concatenate solids into one big graph with offset coedge/face/edge ids.""" - fo = co = eo = 0 - out = {k: [] for k in ["ff", "ef", "fg", "eg", "c2f", "c2e", "c2n", "c2p", "c2m", "rev", "labels"]} - for r in batch: - out["ff"].append(r["ff"]); out["ef"].append(r["ef"]) - out["fg"].append(r["fg"]); out["eg"].append(r["eg"]) - out["c2f"].append(r["c2f"] + fo); out["c2e"].append(r["c2e"] + eo) - out["c2n"].append(r["c2n"] + co); out["c2p"].append(r["c2p"] + co); out["c2m"].append(r["c2m"] + co) - out["rev"].append(r["rev"]); out["labels"].append(r["labels"]) - fo += r["nf"]; eo += r["ef"].shape[0]; co += r["c2f"].shape[0] - cat = {k: np.concatenate(v, axis=0) for k, v in out.items()} - cat["n_faces"] = fo - return cat - - -# --- model ------------------------------------------------------------------ +# --- 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.c1 = mlxnn.Conv2d(cin, 32, 3, padding=1) - self.c2 = mlxnn.Conv2d(32, 64, 3, padding=1) - self.lin = mlxnn.Linear(64, out) + 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): # [F,10,10,7] - x = mlxnn.relu(self.c1(g)); x = mlxnn.relu(self.c2(x)) - x = x.mean(axis=(1, 2)) # global avg pool -> [F,64] - return mlxnn.relu(self.lin(x)) + 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.c1 = mlxnn.Conv1d(cin, 32, 3, padding=1) - self.c2 = mlxnn.Conv1d(32, 64, 3, padding=1) - self.lin = mlxnn.Linear(64, out) + 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): # [E,10,6] - x = mlxnn.relu(self.c1(g)); x = mlxnn.relu(self.c2(x)) - x = x.mean(axis=1) # [E,64] - return mlxnn.relu(self.lin(x)) + 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.ws = mlxnn.Linear(d, d); self.wn = mlxnn.Linear(d, d) - self.wp = mlxnn.Linear(d, d); self.wm = mlxnn.Linear(d, d) + 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.ws(h) + self.wn(h[nidx]) + self.wp(h[pidx]) + self.wm(h[midx])) + return mlxnn.relu(self.W_self(h) + self.W_next(h[nidx]) + self.W_prev(h[pidx]) + self.W_mate(h[midx])) -class BRepNetMLX(mlxnn.Module): - def __init__(self, num_classes=8, ent=64, hid=128, layers=6): +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.surf = SurfEnc(7, 64) - self.curve = CurveEnc(6, 64) - self.face_proj = mlxnn.Linear(8 + 64, ent) - self.edge_proj = mlxnn.Linear(7 + 64, ent) - self.in_proj = mlxnn.Linear(2 * ent + 1, hid) - self.convs = [CoedgeConv(hid) for _ in range(layers)] - self.head = mlxnn.Linear(hid, num_classes) + 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.surf(b["fg"])], axis=1))) - edge_repr = mlxnn.relu(self.edge_proj(mx.concatenate([b["ef"], self.curve(b["eg"])], axis=1))) - h = mx.concatenate([face_repr[b["c2f"]], edge_repr[b["c2e"]], b["rev"][:, None]], axis=1) - h = mlxnn.relu(self.in_proj(h)) - for conv in self.convs: - h = conv(h, b["c2n"], b["c2p"], b["c2m"]) - # coedge -> face mean pool via one-hot matmul (scatter-mean) - nf = b["n_faces"] - onehot = (b["c2f"][:, None] == mx.arange(nf)[None, :]).astype(h.dtype) # [C, F] - face_emb = (onehot.T @ h) / mx.maximum(onehot.sum(axis=0)[:, None], 1) - return self.head(face_emb) # [F, num_classes] - - -def to_mx(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(b["ff"]), "ef": mx.array(b["ef"]), - "fg": mx.array(b["fg"]), "eg": mx.array(b["eg"]), - "c2f": mx.array(b["c2f"]), "c2e": mx.array(b["c2e"]), - "c2n": mx.array(b["c2n"]), "c2p": mx.array(b["c2p"]), "c2m": mx.array(b["c2m"]), - "rev": mx.array(b["rev"]), "n_faces": int(b["n_faces"]), + "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]), } @@ -172,87 +244,86 @@ def miou(conf, nc): def main(): ap = argparse.ArgumentParser() - ap.add_argument("--mode", choices=["probe", "train"], default="train") - ap.add_argument("--n-train", type=int, default=6000) - ap.add_argument("--n-val", type=int, default=1500) - ap.add_argument("--epochs", type=int, default=20) - ap.add_argument("--bs", type=int, default=16) - ap.add_argument("--lr", type=float, default=1e-3) + 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 = json.load(open(_DATA / "dataset.json")) - nc = int(ds["num_classes"]) - fs = ds["feature_standardization"] - fmean = np.array([e["mean"] for e in fs["face_features"]], np.float32) - fstd = np.array([e["standard_deviation"] for e in fs["face_features"]], np.float32) + 1e-8 - emean = np.array([e["mean"] for e in fs["edge_features"]], np.float32) - estd = np.array([e["standard_deviation"] for e in fs["edge_features"]], np.float32) + 1e-8 - print(f"num_classes={nc} classes={ds['class_names']}", flush=True) - - train_stems = ds["training_set"][: args.n_train] - val_stems = ds["validation_set"][: args.n_val] - print(f"loading {len(train_stems)} train / {len(val_stems)} val records ...", flush=True) - train = load_records(train_stems, fmean, fstd, emean, estd) - val = load_records(val_stems, fmean, fstd, emean, estd) - print(f"loaded {len(train)} train / {len(val)} val solids", flush=True) + ds_meta = json.load(open(_DATA / "dataset.json")) + nc = int(ds_meta["num_classes"]) - model = BRepNetMLX(num_classes=nc) + model = LLBRepNetMLX(nc=nc) if args.mode == "probe": - b = to_mx(collate(train[:4])) + 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: batch faces={b['n_faces']} logits={out.shape} finite={bool(mx.isfinite(out).all().item())}", flush=True) + print(f"probe: logits {out.shape} finite={bool(mx.isfinite(out).all().item())}", flush=True) return - opt = optim.Adam(learning_rate=args.lr) - - def loss_fn(b, labels): - logits = model(b) - mask = labels != IGNORE - ce = mlxnn.losses.cross_entropy(logits, mx.where(mask, labels, 0), reduction="none") * mask - return ce.sum() / mx.maximum(mask.sum(), 1) - - lg = mlxnn.value_and_grad(model, loss_fn) - - def evaluate(recs): - conf = np.zeros((nc, nc), np.int64) - for k in range(0, len(recs), args.bs): - b = collate(recs[k:k + args.bs]) - logits = model(to_mx(b)) - pred = np.array(mx.argmax(logits, axis=1).tolist()) - lab = b["labels"] - m = lab != IGNORE - for p, t in zip(pred[m], lab[m]): - conf[t, p] += 1 - acc = float(np.trace(conf) / max(conf.sum(), 1)) - return miou(conf, nc), acc - - best = -1.0 - for epoch in range(args.epochs): - np.random.shuffle(train) - tot = 0.0; nb = 0 - for k in range(0, len(train), args.bs): - b = collate(train[k:k + args.bs]) - labels = mx.array(b["labels"]) - lv, g = lg(to_mx(b), labels) - opt.update(model, g) - mx.eval(model.parameters(), opt.state, lv) - tot += float(lv.item()); nb += 1 - vmiou, vacc = evaluate(val) - print(f"epoch {epoch+1}/{args.epochs} loss={tot/max(nb,1):.4f} val_mIoU={vmiou:.3f} val_acc={vacc:.3f}", flush=True) - if vmiou > best: - best = vmiou - mx.save_safetensors(f"{args.out}/brepnet_mlx.safetensors", dict(tree_flatten(model.parameters()))) - - result = {"framework": "MLX", "task": "B-rep face segmentation", "dataset": "Fusion360", - "num_classes": nc, "n_train": len(train), "n_val": len(val), "epochs": args.epochs, - "best_val_mIoU": round(best, 3), "pytorch_reference_test_mIoU": 0.828, - "checkpoint": f"{args.out}/brepnet_mlx.safetensors"} + 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_DONE", json.dumps(result), flush=True) + print("BREPNET_MLX_PARITY", json.dumps(result), flush=True) if __name__ == "__main__": From b2bdacc75be0705c6cb1b756a36bc16042f8238b Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 07:00:46 -0500 Subject: [PATCH 12/17] feat(ll_ocadr): faithful MLX PointNet++/Point-BERT tower + forward-parity proof The configured ll_ocadr model was never trained (no checkpoint to convert) and can't run on Apple Silicon, so unlike stepnet/brepnet there are no weights to load. But the encoders are plain nn.Modules with no LLM dependency, so faithfulness is still PROVABLE: random-init the real PyTorch GeometryNet (PointNet++) + ShapeNet (Point-BERT, 0.5B config depth=4/heads=8) + linear projector, convert weights into this MLX tower, feed the same (coords, normals) through both, assert forward parity. The 1x1 convs are per-point Linears over channels -> implemented as MLX Linear + BatchNorm (no NHWC Conv layout); MultiheadAttention reconstructed manually from the packed in_proj_weight; FPS + ball-query computed once in numpy (deterministic start) and shared by both models so sampling can't cause divergence. Result: max-abs-diff geometry=1.2e-6, shape=4.1e-6, mesh_tokens=3.3e-6 -> PASS. The MLX tower provably reproduces the real encoders. (Retrain wiring next.) Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_ocadr/mlx/faithful_tower_mlx.py | 415 +++++++++++++++++++++++++++++ 1 file changed, 415 insertions(+) create mode 100644 ll_ocadr/mlx/faithful_tower_mlx.py 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() From 7c4fa7487c82e343b6119066ea8685eb62345a41 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 07:42:31 -0500 Subject: [PATCH 13/17] =?UTF-8?q?feat(ll=5Fgen):=20faithful=20MLX=20port?= =?UTF-8?q?=20of=20STEPVAE=20=E2=80=94=20convert=20real=20weights,=20prove?= =?UTF-8?q?=20parity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ll_gen's neural generator is ll_stepnet.stepnet.vae.STEPVAE (transformer encoder-decoder VAE over CAD command tokens), trained to vae_warm.pt / vae_rl_solid.pt. This reproduces the EXACT architecture in MLX and converts the real weights: encoder (6 post-norm layers, masked-mean pool) + mu/log_var heads; decode(z) = latent_project broadcast + dec_pos -> 6x TransformerDecoderLayer (causal self-attn + cross-attn to z-memory + FFN, 3 post-norm LayerNorms) -> command_head(6) + 16 param_heads(256). MHA splits the packed in_proj into Wq/Wk/Wv so one module serves both self- and cross-attention; the parallel decode means decoder.token/pos embeddings are unused on this path (skipped). Verified by converting vae_warm.pt (227 tensors) and comparing the full encode-> decode against the PyTorch STEPVAE on the same tokens: max-abs-diff mu=4e-7, command_logits=1.6e-6, param_logits=1.1e-5, decoded-command argmax agreement = 1.0 -> PASS. The MLX VAE IS the trained generator, running natively on Apple Silicon. Modes: probe | convert | parity. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/mlx/vae_mlx.py | 301 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 301 insertions(+) create mode 100644 ll_gen/mlx/vae_mlx.py 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() From 6b9071abb8a91459a65f190f3cf536cf43ecb886 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 09:07:38 -0500 Subject: [PATCH 14/17] feat(ll_ocadr): retrain on the faithful PointNet++/Point-BERT tower (256 mesh tokens) Replaces the maxpool stand-in encoder with the parity-proven faithful tower (faithful_tower_mlx.py): real GeometryNet (PointNet++) + ShapeNet (Point-BERT) + linear projector -> 256 mesh tokens spliced into a frozen 4-bit Qwen2; tower + LoRA train jointly (31.98M trainable). Data extended to N=2048 points + area-weighted normals + cached deterministic FPS/ball-query grouping. Task/data/metrics held identical to the maxpool run for a direct comparison. Honest result (balanced 3-way class, 1230 train / 246 val): llm_generation_acc = 0.919 vs shuffled-mesh baseline = 0.313 (majority 0.374) encoder_mesh_read_acc (aux head) = 0.785 The 0.61 grounding gap (vs the maxpool encoder's 0.22: 0.543 vs 0.322) shows the real architecture gives the LLM far stronger geometric signal. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_ocadr/mlx/train_ocadr_mlx.py | 297 ++++++++++++++++---------------- 1 file changed, 153 insertions(+), 144 deletions(-) diff --git a/ll_ocadr/mlx/train_ocadr_mlx.py b/ll_ocadr/mlx/train_ocadr_mlx.py index eea4b69..7dc5a4b 100644 --- a/ll_ocadr/mlx/train_ocadr_mlx.py +++ b/ll_ocadr/mlx/train_ocadr_mlx.py @@ -1,19 +1,18 @@ -"""ll_ocadr in MLX, end to end on Apple Silicon. - -A native-MLX implementation of the ll_ocadr multimodal idea: a trainable 3D -geometry encoder + projector map a CAD point cloud into the embedding space of a -4-bit-quantized Qwen2 LLM (via mlx-lm's ``input_embeddings`` path). The LLM base -weights stay frozen/quantized, but LoRA adapters on its attention/MLP are trained -jointly with the encoder so the LLM actually LEARNS TO ATTEND to the injected -mesh tokens (LLaVA-style; a frozen LLM otherwise ignores the out-of-distribution -mesh embeddings and collapses to the text prior). This trains on Apple Silicon -where PyTorch-MPS would OOM, and the projector/LoRA are sized to the real LLM. - -Task (mesh-grounded): given a CAD point cloud, generate a short structured -description "faces= ext= class=" derived from the -actual solid. Honest success bar: held-out face-count-class accuracy must beat a -SHUFFLED-mesh baseline (same model, mesh swapped) — i.e. the model must read the -geometry, not the text prior. +"""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. """ @@ -51,6 +50,9 @@ _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 @@ -59,7 +61,7 @@ from OCC.Core.BRep import BRep_Tool # noqa: E402 from OCC.Core.TopLoc import TopLoc_Location # noqa: E402 -NUM_POINTS, NUM_MESH_TOKENS = 512, 16 +NUM_POINTS, NUM_MESH_TOKENS = 2048, 256 CLASS_NAMES = ["simple", "box", "complex"] @@ -67,9 +69,20 @@ def _bucket(nf: int) -> int: return 0 if nf <= 4 else (1 if nf <= 6 else 2) -def _sample_points(shape, P=NUM_POINTS): +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) - pts = [] + verts, tris = [], [] exp = TopExp_Explorer(shape, TopAbs_FACE) nfaces = 0 while exp.More(): @@ -79,90 +92,80 @@ def _sample_points(shape, P=NUM_POINTS): 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) - pts.append([p.X(), p.Y(), p.Z()]) + 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 not pts: - return None, 0 - pts = np.asarray(pts, np.float32) - c = pts.mean(0) - s = np.abs(pts - c).max() + 1e-6 - pts = (pts - c) / s - idx = np.random.choice(len(pts), P, replace=len(pts) < P) - return pts[idx], nfaces + 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, allow_pickle=True) - if int(d["clouds"].shape[0]) >= n_target: - return d["clouds"][:n_target], list(d["texts"][:n_target]), d["classes"][:n_target] + 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"))) - clouds, texts, classes = [], [], [] + coords, normals, sa1, sa2xyz, sa2idx, classes = [], [], [], [], [], [] for f in files: - if len(clouds) >= n_target: + 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) - ne = len(cad.seq) shape = execute(command_dicts(cad)) if shape is None: continue - pc, nf = _sample_points(shape) + pc, nrm, nf = _sample_points(shape) if pc is None or nf < 1: continue - clouds.append(pc) - texts.append(f"faces={nf} ext={ne} class={CLASS_NAMES[_bucket(nf)]}") + 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 - clouds = np.stack(clouds) - classes = np.array(classes) + 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, clouds=clouds, texts=np.array(texts, dtype=object), classes=classes) - return clouds, texts, classes - - -class GeometryEncoderMLX(mlxnn.Module): - """PointNet-style encoder: per-point MLP + MAXPOOL global descriptor (the - proven-discriminative path — an encoder-only maxpool head reaches ~0.87 on - this task), expanded into NUM_MESH_TOKENS distinct tokens in the LLM hidden - space. Every mesh token is a learned projection of the discriminative global - feature, so both the aux head and the (LoRA) LLM get strong signal.""" - - def __init__(self, d=256, llm_hidden=896, n_tokens=NUM_MESH_TOKENS): - super().__init__() - self.point_mlp = mlxnn.Sequential( - mlxnn.Linear(3, d), mlxnn.GELU(), - mlxnn.Linear(d, d), mlxnn.GELU(), - mlxnn.Linear(d, d), mlxnn.GELU(), - ) - self.n_tokens = n_tokens - self.llm_hidden = llm_hidden - self.token_proj = mlxnn.Linear(d, n_tokens * llm_hidden) # global -> N tokens - self.aux_head = mlxnn.Linear(llm_hidden, 3) - - def __call__(self, points): - feats = self.point_mlp(points) # [B, P, d] - glob = feats.max(axis=1) # [B, d] maxpool (discriminative) - B = glob.shape[0] - return self.token_proj(glob).reshape(B, self.n_tokens, self.llm_hidden) - - def aux_logits(self, mesh_tokens): - return self.aux_head(mesh_tokens.mean(axis=1)) # [B, 3] + np.savez(cache, **out) + return out class OCADRMLX(mlxnn.Module): - """Wraps the trainable encoder + the LoRA-adapted (otherwise frozen) LLM so - mlx value_and_grad differentiates encoder params AND LoRA params together.""" + """Trainable faithful tower + LoRA-adapted (otherwise frozen) LLM, so mlx + value_and_grad differentiates tower params AND LoRA params together.""" - def __init__(self, encoder, llm): + def __init__(self, tower, llm): super().__init__() - self.encoder = encoder + self.tower = tower self.llm = llm @@ -170,10 +173,10 @@ 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=3000) - ap.add_argument("--n-val", type=int, default=600) - ap.add_argument("--epochs", type=int, default=12) - ap.add_argument("--bs", type=int, default=16) + 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) @@ -191,23 +194,28 @@ def command_dicts(cad): quantization_bits=8, normalization_range=2.0) deps = (CADSequence, command_dicts, execute_command_proposal) - print(f"loading {args.llm} + applying LoRA (rank={args.lora_rank}, layers={args.lora_layers}) ...", flush=True) + 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}) + linear_to_lora_layers(llm, args.lora_layers, {"rank": args.lora_rank, "scale": 20.0, "dropout": 0.0}) H = llm.args.hidden_size - enc = GeometryEncoderMLX(llm_hidden=H) - model = OCADRMLX(enc, llm) + 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 params (encoder + LoRA): {n_trainable/1e6:.2f}M; LLM base frozen+quantized", flush=True) + 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 make_batch(clouds, texts): - mesh = model.encoder(mx.array(clouds)) # [B,T,H] - prompt_emb = model.llm.model.embed_tokens(prompt_ids) # [Lp,H] + 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] @@ -224,58 +232,57 @@ def make_batch(clouds, texts): labels.append([-100] * (Lp + T) + r.tolist() + [-100] * pad) return mx.stack(seqs), mx.array(np.array(labels)), mesh - def loss_fn(clouds, texts, classes): - emb, labels, mesh = make_batch(clouds, texts) + 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) - # auxiliary classification on the mesh tokens (makes them discriminative) - aux = mlxnn.losses.cross_entropy(model.encoder.aux_logits(mesh), mx.array(classes), reduction="mean") + 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": - clouds, texts, classes = build_dataset(8, deps, None) - lv, grads = mlxnn.value_and_grad(model, loss_fn)(clouds[:4], texts[:4], classes[:4]) + 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} trainable grad_norm={gnorm:.4f} text={texts[0]!r}", flush=True) + print(f"probe loss={float(lv.item()):.4f} grad_norm={gnorm:.4f}", flush=True) return - cache = f"{args.out}/ocadr_data_cache.npz" os.makedirs(args.out, exist_ok=True) - print("building/loading dataset ...", flush=True) - allc, _allt, allcl = build_dataset(args.n_train + args.n_val, deps, cache) - # Balance the 3 classes so the text prior gives only chance (1/3): this - # removes the majority-class shortcut and forces the model to READ the mesh - # to beat the shuffled-mesh baseline. Target is the bare class word (no - # predictable "faces=N" prefix to ride). + 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) - allc, allcl = allc[keep], allcl[keep] - allt = [CLASS_NAMES[c] for c in allcl] - nval = max(len(keep) // 6, 90) - vc, vt, vcl = allc[:nval], allt[:nval], allcl[:nval] - tc, tt, tcl = allc[nval:], allt[nval:], allcl[nval:] - print(f"balanced {tc.shape[0]} train / {vc.shape[0]} val; per-class={per}; target=class-only; " - f"val dist={np.bincount(vcl, minlength=3).tolist()} (majority baseline={np.bincount(vcl,minlength=3).max()/len(vcl):.3f})", - flush=True) + 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) - n = tc.shape[0] - def greedy_class(clouds): - B = clouds.shape[0] - mesh = model.encoder(mx.array(clouds)) + 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(12): + 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) @@ -285,45 +292,47 @@ def greedy_class(clouds): preds = np.full(B, -1, int) for b in range(B): txt = tok.decode(toks[b]).lower() - # target is the bare class word; take the first class name that appears 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(clouds, classes, shuffle=False): - src = np.random.permutation(clouds.shape[0]) if shuffle else np.arange(clouds.shape[0]) - cl = clouds[src] + 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, clouds.shape[0], 32): - correct += int((greedy_class(cl[k:k + 32]) == classes[k:k + 32]).sum()) - return correct / clouds.shape[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(clouds, classes): - """Accuracy of the encoder's auxiliary head — does the MLX geometry - encoder itself read the mesh (independent of LLM verbalization)?""" + def aux_acc(d, idx, classes): + model.tower.eval() correct = 0 - for k in range(0, clouds.shape[0], 64): - mesh = model.encoder(mx.array(clouds[k:k + 64])) - pred = np.array(mx.argmax(model.encoder.aux_logits(mesh), axis=1).tolist()) - correct += int((pred == classes[k:k + 64]).sum()) - return correct / clouds.shape[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): - idx = perm[k:k + args.bs] - lv, grads = lg_fn(tc[idx], [tt[i] for i in idx], tcl[idx]) + 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(vc, vcl) - shuf = class_acc(vc, vcl, shuffle=True) - aux = aux_acc(vc, vcl) + 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: @@ -331,19 +340,19 @@ def aux_acc(clouds, classes): mx.save_safetensors(f"{args.out}/ocadr_mlx.safetensors", dict(tree_flatten(model.trainable_parameters()))) - majority = float(np.bincount(vcl, minlength=3).max() / len(vcl)) - result = {"framework": "MLX", "task": "CAD point-cloud -> class (simple/box/complex)", - "llm": args.llm, "llm_base": "frozen+4bit-quantized", - "trainable": "geometry encoder + LoRA adapters", "trainable_params_M": round(n_trainable / 1e6, 2), - "n_train": int(tc.shape[0]), "n_val": int(vc.shape[0]), "epochs": args.epochs, - "majority_baseline": round(majority, 3), - "encoder_mesh_read_acc": round(aux_acc(vc, vcl), 3), - "llm_generation_acc": round(class_acc(vc, vcl), 3), - "shuffled_mesh_baseline": round(class_acc(vc, vcl, shuffle=True), 3), + 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_DONE", json.dumps(result), flush=True) + print("OCADR_MLX_FAITHFUL_DONE", json.dumps(result), flush=True) if __name__ == "__main__": From a9b0bdc48fdc04ba2fcf945de9ff61f231e0e7d8 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 10:12:25 -0500 Subject: [PATCH 15/17] =?UTF-8?q?feat(ll=5Fgen):=20autoregressive=20CAD-co?= =?UTF-8?q?mmand=20generator=20=E2=80=94=20measured-valid=20CAD=20(0.91+)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing generators cannot produce valid CAD: the command-VAE's parallel z-broadcast decoder is primitive-limited / posterior-collapses (~0-12% valid) and the diffusion path samples faces independently so they never mate (0% valid). This adds the robust, proven route (DeepCAD/Text2CAD): generate the construction PROGRAM autoregressively and execute it, so the OCC kernel builds the solid command-by-command. A causal-transformer LM over the 268-token command vocabulary is teacher-forced on 38k real DeepCAD command sequences, then sampled autoregressively (temperature 1.0, top-k 20) -> decoded to command_dicts -> executed through the REAL OCC kernel. Validity is MEASURED and gated honestly against the cylinder trap: a sample counts only if it forms a solid (solid_count>=1) with non-degenerate volume (>1e-4); we also report num_distinct (rounded bounding boxes), top_shape_frac (mode-collapse guard), and the volume spread. Result (256 samples, real kernel): validity 0.914 (234/256), num_distinct 104, mean_vol 1.84, vol p10-p90 [0.15, 3.64] best 0.969 vs untrained baseline 0.023 (and vs VAE ~0.12 / diffusion 0.0). High distinct + wide volume spread => genuinely diverse, non-degenerate solids, not one repeated trivial shape. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/mlx/ar_generator_mlx.py | 404 +++++++++++++++++++++++++++++++++ 1 file changed, 404 insertions(+) create mode 100644 ll_gen/mlx/ar_generator_mlx.py 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() From a68606ec6169cfb581978b1a32a9f47be4cbc251 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 10:42:12 -0500 Subject: [PATCH 16/17] fix(ll_gen): honest validity gate + diffusion sampling device fix Two coupled fixes surfaced while making the generators produce *measured*-valid CAD: 1. Deceptive validity metric. compute_validity_rate counted r.is_valid (BRepCheck) only, which passes volume-less shells and zero-volume degenerates. An independently-sampled-face generator therefore scored validity=1.0 while producing ZERO real solids. Added GenerationMetrics.is_valid_solid (requires a closed solid with positive volume when a geometry report exists; falls back to is_valid for abstract test stand-ins) and routed validity_rate, num_valid and distinct_valid through it. Updated the test fixture's "valid" result to carry a real volume. 2. Diffusion sampling cpu/mps crash. StructuredDiffusion.sample() honored a caller -supplied device that could disagree with the denoiser weights (a generator whose .device attr was not updated after the model was .to()'d), crashing at the first Linear ("input is on cpu but expected on mps"). sample() now always computes on the weights' own device. Effect: the diffusion eval RUNS now instead of crashing, and the honest metric reports its true validity=0.000 (its independently-sampled faces never mate into a solid) instead of the former hollow 1.000. Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/ll_gen/training/metrics.py | 29 +++++++++++++++++++++++++---- ll_gen/tests/test_validity_eval.py | 7 ++++--- ll_stepnet/stepnet/diffusion.py | 10 ++++++++-- 3 files changed, 37 insertions(+), 9 deletions(-) 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/tests/test_validity_eval.py b/ll_gen/tests/test_validity_eval.py index 0447e71..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), ) diff --git a/ll_stepnet/stepnet/diffusion.py b/ll_stepnet/stepnet/diffusion.py index 9bbfb7a..c18888b 100644 --- a/ll_stepnet/stepnet/diffusion.py +++ b/ll_stepnet/stepnet/diffusion.py @@ -848,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 From c1e261b159ad05f1858f34d059df989389392ed9 Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Thu, 11 Jun 2026 11:48:49 -0500 Subject: [PATCH 17/17] =?UTF-8?q?feat(ll=5Fgen):=20latent=20diffusion=20th?= =?UTF-8?q?at=20generates=20valid=20CAD=20=E2=80=94=20the=20real=20diffusi?= =?UTF-8?q?on=20fix?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shipped StructuredDiffusion denoises raw B-rep geometry (independent face UV-grids + edge polylines) then tries to sew them; independently-sampled faces never mate, so honest validity is 0. This replaces that with DeepCAD's actual generative design: diffuse the CONSTRUCTION-PROGRAM latent and decode with an execution-respecting autoregressive decoder, so the OCC kernel builds a watertight solid. - SeqAutoencoder: deterministic (not a VAE -> no posterior collapse), encoder -> z[64]; z-conditioned causal AR decoder trained teacher-forced with 50% word-dropout so the latent must carry the global program. - LatentDDPM: eps-prediction denoiser over the normalised z bank. Honest acceptance bar = validity of DIFFUSION-SAMPLED z (noise -> denoise -> decode -> execute), real OCC kernel, solid+volume gated, distinct>1, beating a z=0 predict-the-mean baseline: sampled-z : validity 0.934 (239/256), distinct 138, top_shape 11.7%, mean_vol 1.32 z=0 base : validity 1.000 but distinct 14, top_shape 47.7% (near mode-collapse) Validity comes from the decoder (both are high); the DIFFUSION's contribution is the DISTRIBUTION — 138 distinct valid shapes vs the baseline's 14. It samples the diverse space of valid programs instead of repeating the mean. (vs the old geometry-diffusion honest validity 0.0.) Co-Authored-By: Claude Opus 4.8 (1M context) --- ll_gen/mlx/latent_diffusion_mlx.py | 359 +++++++++++++++++++++++++++++ 1 file changed, 359 insertions(+) create mode 100644 ll_gen/mlx/latent_diffusion_mlx.py 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()