From f5a2719f7a1c5341e22ad0c0ea3c2082859dadfe Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 17:37:53 +0100 Subject: [PATCH] Split optimizer preparation from solving --- README.md | 1 + docs/implementation-plan.md | 13 +++ docs/optimizer-algorithm-map.md | 18 ++- src/repixelizer/continuous.py | 200 +++++++++++++++++++++++--------- 4 files changed, 169 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 4a3c2e7..84dc611 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,7 @@ Current tile-graph status: - the latest reduce-by-key pinned badge probe under `artifacts/tile-graph-reducebykey-v1-badge-126/` finishes in about `16.1s` cold and `6.0s` warm on CUDA, versus the previous full-resolution cut at about `96.9s` cold and `5.4s` warm - that speed pass is a real throughput win but a small quality regression: the same pinned badge case moves from `0.1814` to `0.2036` source-fidelity - the optimizer now has its own deep-dive map in `docs/optimizer-algorithm-map.md`, and the first cut from that map removed the old UV-optimizer name plus the continuous k-means cluster-boundary scout; the continuous path is now a fixed-lattice, source-first discrete chooser guided by real edge evidence +- optimizer preparation is now split into a private `_OptimizerPrep` bundle, so the continuous path has a visible seam between "build the map" and "choose the pixels" - after the first pruning pass, tile-graph no longer participates in pipeline phase-rerank probes and no longer carries hybrid geometry priors through its unary cost; the path is now one lattice-conditioned candidate generator plus one local discrete solver - this fixes the core design mismatch that had allowed repeated distant labels to create big same-color patches and opaque black background blocks - on the current `24x24` emblem smoke case, an end-to-end `tile-graph` run dropped from about `2.57s` on CPU to `0.61s` on CUDA on this machine diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 1a96508..578ab1f 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -121,6 +121,19 @@ The first optimizer cut landed immediately after that map: - `continuous.py` now uses only real source edges for guidance instead of blending in coarse cluster boundaries - the pinned cleaned-badge continuous smoke run under `artifacts/optimizer-cut-v1-badge-126/` still behaves coherently after the cut: snap scores `0.08100`, final scores `0.07535`, and the generated preview is `diagnostics/output-preview.png` +The second optimizer cut is structural: + +- `_prepare_optimizer(...)` now owns source tensor creation, UV construction, edge gradients, source lattice reference building, detail reference building, source reliability, source deltas, guidance, and cell geometry +- `_source_detail_delta_tensors(...)` computes premultiplied source-detail deltas once instead of repeating the same `premultiply(...)` calls for each axis +- `optimize_lattice_pixels(...)` is now an orchestration shell around prep, snap, refine, the source-fidelity guardrail, and artifact packaging +- the pinned cleaned-badge continuous smoke run under `artifacts/optimizer-cut-v2-prep-badge-126/` is byte-size identical to the previous preview and preserves the same source-fidelity numbers: snap `0.08100`, final `0.07535` + +Next after that: + +- put the representative portrait on trial in refine +- add a focused ablation runner for pinned badge/emblem cases before cutting more behavior +- collapse duplicated adjacency / motif / line voices only after ablation shows which stage owns each idea + The next optimizer simplification pass should use that map as its cutting checklist instead of guessing from scattered helper names. ## High-value regression case diff --git a/docs/optimizer-algorithm-map.md b/docs/optimizer-algorithm-map.md index 6417948..56f5871 100644 --- a/docs/optimizer-algorithm-map.md +++ b/docs/optimizer-algorithm-map.md @@ -40,6 +40,8 @@ The optimizer carries five important pieces of cargo all the way through: If we lose track of one of those, the optimizer becomes impossible to reason about. +The first structural cut now gathers that cargo into `_OptimizerPrep` before snap and refine begin. That does not make the algorithm better by itself, but it makes the machine easier to see: preparation builds the map, snap and refine walk it. + ## Stage 0: The pipeline picks the ruler Functions: @@ -220,6 +222,12 @@ The optimizer now has two different stories about the same lattice: Much of the optimizer's complexity is the cost of arbitrating between those two stories. +Current code shape: + +- `_prepare_optimizer(...)` owns this whole setup stage now +- `_source_detail_delta_tensors(...)` computes the source-detail deltas once instead of rebuilding premultiplied references several times +- `optimize_lattice_pixels(...)` now receives a prepared bundle and orchestrates snap, refine, and the final guardrail + ## Stage 5: Snap every output cell to one real nearby source pixel Function: @@ -468,7 +476,7 @@ These are the places where the optimizer still speaks in an overly complicated o - the machine still maintains two overlapping portraits of the same lattice and spends much of its complexity mediating the argument between them - snap, relax, and refine all carry slightly different versions of adjacency, motif, line, and delta agreement; the same idea is being said three times with different accents - the candidate generator is local and the solver is discrete, but the vocabulary around it still sounds like a continuous deformation engine -- the main entry point is now honestly named `optimize_lattice_pixels(...)`, and the k-means boundary scout is gone; those cuts removed two contradictions from the previous map +- the main entry point is now honestly named `optimize_lattice_pixels(...)`, the k-means boundary scout is gone, and prep is split from decision-making; those cuts removed three contradictions from the previous map This is probably the real cutting checklist for the next optimizer pass. @@ -484,11 +492,11 @@ That is much simpler than the surrounding names make it sound. If we keep pruning, the highest-value cuts look like this: -1. Separate preparation from decision-making. - The edge scout, portraits, reliability maps, and delta maps want to be a clear prep stage instead of being braided through one giant function. -2. Decide whether the representative portrait still earns its keep everywhere. +1. Decide whether the representative portrait still earns its keep everywhere. If the source-first path is now strong enough, some of that arbitration machinery may be removable. -3. Collapse duplicate structure voices. +2. Collapse duplicate structure voices. Adjacency, motif, and line are important, but they should not need three near-parallel dialects unless they are truly doing different work. +3. Put relax on trial. + The soft relaxation stage may be a useful bridge, or it may be an expensive second solver that mostly repeats refine. That is the next place to swing the axe. diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index 6bee610..4cfe95d 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -1,5 +1,7 @@ from __future__ import annotations +from dataclasses import dataclass + import numpy as np from .io import premultiply, unpremultiply @@ -9,6 +11,23 @@ from .source_reference import build_source_lattice_reference from .types import ContinuousSourceAnalysis, InferenceResult, SolverArtifacts +@dataclass(slots=True) +class _OptimizerPrep: + source_t: object + uv_t: object + representative_t: object + source_reference_t: object + source_reliability_t: object + source_lattice_reference: object + source_delta_x_t: object | None + source_delta_y_t: object | None + source_delta_diag_t: object | None + source_delta_anti_t: object | None + guidance: np.ndarray + cell_x: float + cell_y: float + + def _require_torch(): try: import torch @@ -426,6 +445,50 @@ def _build_source_detail_reference(source_rgba: np.ndarray, source_reference, so return detail_rgba.astype(np.float32) +def _source_detail_delta_tensors( + torch, + source_detail_reference: np.ndarray, + *, + target_width: int, + target_height: int, + device: str, +): + detail_premul = premultiply(source_detail_reference) + source_delta_x_t = ( + torch.from_numpy((detail_premul[:, 1:, :] - detail_premul[:, :-1, :])[None, ...]).to( + device=device, + dtype=torch.float32, + ) + if target_width > 1 + else None + ) + source_delta_y_t = ( + torch.from_numpy((detail_premul[1:, :, :] - detail_premul[:-1, :, :])[None, ...]).to( + device=device, + dtype=torch.float32, + ) + if target_height > 1 + else None + ) + source_delta_diag_t = ( + torch.from_numpy((detail_premul[1:, 1:, :] - detail_premul[:-1, :-1, :])[None, ...]).to( + device=device, + dtype=torch.float32, + ) + if target_height > 1 and target_width > 1 + else None + ) + source_delta_anti_t = ( + torch.from_numpy((detail_premul[1:, :-1, :] - detail_premul[:-1, 1:, :])[None, ...]).to( + device=device, + dtype=torch.float32, + ) + if target_height > 1 and target_width > 1 + else None + ) + return source_delta_x_t, source_delta_y_t, source_delta_diag_t, source_delta_anti_t + + def _edge_reliability(source_reliability, axis: str): if axis == "x": return (source_reliability[:, 1:] + source_reliability[:, :-1]) * 0.5 @@ -1336,21 +1399,16 @@ def _discrete_refine_output( return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history -def optimize_lattice_pixels( +def _prepare_optimizer( + torch, + F, rgba: np.ndarray, inference: InferenceResult, analysis: ContinuousSourceAnalysis, - steps: int, - seed: int, + solver_params: SolverHyperParams, + *, device: str, - solver_params: SolverHyperParams | None = None, -) -> SolverArtifacts: - torch, F = _require_torch() - device = _resolve_device(torch, device) - solver_params = solver_params or SolverHyperParams() - torch.manual_seed(seed) - if device == "cuda": - torch.cuda.manual_seed_all(seed) +) -> _OptimizerPrep: source = premultiply(rgba) height, width = source.shape[:2] cell_x = width / max(1, inference.target_width) @@ -1359,7 +1417,7 @@ def optimize_lattice_pixels( source_edge = analysis.edge_map.astype(np.float32) edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge) edge_t = torch.from_numpy(source_edge[None, None, ...]).to(device=device, dtype=torch.float32) - uv0 = _make_regular_uv( + uv = _make_regular_uv( height=height, width=width, target_height=inference.target_height, @@ -1367,7 +1425,7 @@ def optimize_lattice_pixels( phase_x=inference.phase_x, phase_y=inference.phase_y, ) - uv0_t = torch.from_numpy(uv0[None, ...]).to(device=device, dtype=torch.float32) + uv_t = torch.from_numpy(uv[None, ...]).to(device=device, dtype=torch.float32) offsets_t = torch.from_numpy( _make_patch_offsets(height=height, width=width, target_height=inference.target_height, target_width=inference.target_width) ).to(device=device, dtype=torch.float32) @@ -1384,10 +1442,10 @@ def optimize_lattice_pixels( edge_grad_y_hint=edge_grad_y, ) source_detail_reference = _build_source_detail_reference(rgba, source_lattice_reference, solver_params) - initial_patches = _sample_cell_patches(F, source_t, uv0_t, offsets_t) - initial_representative_t, _ = _representative_colors(initial_patches, solver_params) - initial_representative_t = initial_representative_t.detach() - initial_source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to( + initial_patches = _sample_cell_patches(F, source_t, uv_t, offsets_t) + representative_t, _ = _representative_colors(initial_patches, solver_params) + representative_t = representative_t.detach() + source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to( device=device, dtype=torch.float32, ) @@ -1395,37 +1453,65 @@ def optimize_lattice_pixels( device=device, dtype=torch.float32, ) - source_delta_x_t = ( - torch.from_numpy((premultiply(source_detail_reference)[:, 1:, :] - premultiply(source_detail_reference)[:, :-1, :])[None, ...]).to(device=device, dtype=torch.float32) - if inference.target_width > 1 - else None + source_delta_x_t, source_delta_y_t, source_delta_diag_t, source_delta_anti_t = _source_detail_delta_tensors( + torch, + source_detail_reference, + target_width=inference.target_width, + target_height=inference.target_height, + device=device, ) - source_delta_y_t = ( - torch.from_numpy((premultiply(source_detail_reference)[1:, :, :] - premultiply(source_detail_reference)[:-1, :, :])[None, ...]).to(device=device, dtype=torch.float32) - if inference.target_height > 1 - else None + return _OptimizerPrep( + source_t=source_t, + uv_t=uv_t, + representative_t=representative_t, + source_reference_t=source_reference_t, + source_reliability_t=source_reliability_t, + source_lattice_reference=source_lattice_reference, + source_delta_x_t=source_delta_x_t, + source_delta_y_t=source_delta_y_t, + source_delta_diag_t=source_delta_diag_t, + source_delta_anti_t=source_delta_anti_t, + guidance=guide_small[0, 0].detach().cpu().numpy().astype(np.float32), + cell_x=cell_x, + cell_y=cell_y, ) - source_delta_diag_t = ( - torch.from_numpy((premultiply(source_detail_reference)[1:, 1:, :] - premultiply(source_detail_reference)[:-1, :-1, :])[None, ...]).to(device=device, dtype=torch.float32) - if inference.target_height > 1 and inference.target_width > 1 - else None - ) - source_delta_anti_t = ( - torch.from_numpy((premultiply(source_detail_reference)[1:, :-1, :] - premultiply(source_detail_reference)[:-1, 1:, :])[None, ...]).to(device=device, dtype=torch.float32) - if inference.target_height > 1 and inference.target_width > 1 - else None + + +def optimize_lattice_pixels( + rgba: np.ndarray, + inference: InferenceResult, + analysis: ContinuousSourceAnalysis, + steps: int, + seed: int, + device: str, + solver_params: SolverHyperParams | None = None, +) -> SolverArtifacts: + torch, F = _require_torch() + device = _resolve_device(torch, device) + solver_params = solver_params or SolverHyperParams() + torch.manual_seed(seed) + if device == "cuda": + torch.cuda.manual_seed_all(seed) + prep = _prepare_optimizer( + torch, + F, + rgba, + inference, + analysis, + solver_params, + device=device, ) snap_rgba = _snap_output_to_source_pixels( torch, - source_t, - uv0_t, - initial_representative_t, - initial_source_reference_t, - source_reliability_t, - source_lattice_reference, + prep.source_t, + prep.uv_t, + prep.representative_t, + prep.source_reference_t, + prep.source_reliability_t, + prep.source_lattice_reference, solver_params, - cell_x=cell_x, - cell_y=cell_y, + cell_x=prep.cell_x, + cell_y=prep.cell_y, ) snap_t = torch.from_numpy(premultiply(snap_rgba).transpose(2, 0, 1)[None, ...]).to( device=device, @@ -1435,20 +1521,20 @@ def optimize_lattice_pixels( target_rgba, loss_history = _discrete_refine_output( torch, F, - source_t, - uv0_t, - initial_representative_t, - initial_source_reference_t, - source_reliability_t, - source_lattice_reference, + prep.source_t, + prep.uv_t, + prep.representative_t, + prep.source_reference_t, + prep.source_reliability_t, + prep.source_lattice_reference, snap_t, - source_delta_x_t, - source_delta_y_t, - source_delta_diag_t, - source_delta_anti_t, + prep.source_delta_x_t, + prep.source_delta_y_t, + prep.source_delta_diag_t, + prep.source_delta_anti_t, solver_params, - cell_x=cell_x, - cell_y=cell_y, + cell_x=prep.cell_x, + cell_y=prep.cell_y, iterations=steps, ) snap_score = source_lattice_consistency_breakdown( @@ -1473,13 +1559,11 @@ def optimize_lattice_pixels( target_rgba = snap_rgba loss_history.append(float(snap_score)) initial_rgba = snap_rgba - uv_np = uv0_t.detach().cpu().numpy()[0] - guidance = guide_small[0, 0].detach().cpu().numpy() return SolverArtifacts( target_rgba=target_rgba, - uv_field=uv_np, - guidance_strength=guidance.astype(np.float32), + uv_field=prep.uv_t.detach().cpu().numpy()[0], + guidance_strength=prep.guidance, initial_rgba=initial_rgba, loss_history=loss_history, )