From 185cbdf7142b7e82289735cb5b0f543c01591567 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 02:00:21 +0100 Subject: [PATCH] Add fixed lattice tile-graph controls --- README.md | 3 + docs/implementation-plan.md | 20 +++--- src/repixelizer/cli.py | 19 ++++++ src/repixelizer/compare.py | 10 +++ src/repixelizer/inference.py | 50 +++++++++++++- src/repixelizer/pipeline.py | 77 ++++++++++++++++++++- src/repixelizer/tile_graph.py | 123 +++++++++++++++++++++++++++++----- tests/test_cli.py | 28 ++++++++ tests/test_inference.py | 35 ++++++++++ tests/test_pipeline.py | 92 +++++++++++++++++++++++++ tests/test_tile_graph.py | 35 +++++++++- 11 files changed, 459 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 62b3a7e..27b1e91 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,8 @@ Current tile-graph status: - profiling the selected badge candidate shows the main remaining bottleneck is model construction, not the solver loop: `build_tile_graph_model(...)` takes about `142.3s`, and about `131.2s` of that is `_extract_source_region_tiles(...)` - the low-confidence badge rerank is also expensive because it still probes eight lattice candidates before the final pick; on the same fixture that phase takes about `223.9s`, of which about `208.1s` is reconstruction work - the remaining large-fixture bottleneck is the per-component one-cell window cutting pass inside source-region extraction, which is still Python/NumPy-heavy after the GPU CCL stage +- the pipeline now has a direct-control path for iteration: `--target-width` / `--target-height` plus optional `--phase-x` / `--phase-y` let you run an exact lattice without paying the full lattice search, and `--skip-phase-rerank` lets you keep the pipeline from second-guessing that choice +- repeated fixed-lattice `tile-graph` runs in the same Python process now reuse the expensive model build; on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, the first CUDA run took about `10.2s` and the second cached rerun took about `2.1s` - elongated source regions now get a stroke-aware slicing pass that follows their principal axis instead of only marching with cardinal queue steps, and the repo now has a shallow-stroke regression that checks for fan-out on that kind of component - that stroke-aware slicer is a real synthetic improvement but not a real-badge win yet: the latest badge probe under `artifacts/badge-tile-graph-stroke-v2-cuda/` still keeps the initial assignment and slightly regresses source-fidelity (`0.1832` vs `0.1800`) - the new low-risk hybrid mode keeps tile-graph's source-owned candidates but scores them against a continuous-optimizer geometry prepass; on the cleaned badge it improves tile-graph from `0.1832` to `0.1785` under `artifacts/badge-hybrid-v2-cuda/` @@ -72,6 +74,7 @@ repixelize input.png --out output.png --diagnostics-dir diagnostics --device aut repixelize input.png --out output.png --reconstruction-mode tile-graph --diagnostics-dir diagnostics --device cpu repixelize input.png --out output.png --reconstruction-mode tile-graph --diagnostics-dir diagnostics --device cuda repixelize input.png --out output.png --reconstruction-mode hybrid --diagnostics-dir diagnostics --device cuda +repixelize input.png --out output.png --reconstruction-mode tile-graph --target-width 126 --target-height 126 --phase-x 0.0 --phase-y -0.2 --skip-phase-rerank --device cuda ``` Run the optimizer plus baselines: diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index bbf0b59..87bc2a0 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -20,7 +20,7 @@ Implemented: Verified: - local editable install in a dedicated venv -- full test suite currently passing (`60 passed`) +- full test suite currently passing (`66 passed`) - compare-mode smoke run against a real emblem image - the cleaned real badge fixture in `tests/fixtures/real/ai-badge-cleaned.png` still beats naive resize on source-lattice consistency on the continuous path - the latest full-emblem tile-graph atomic probe under `artifacts/full-emblem-tile-graph-atomic-v3-cuda/` lands at `0.0224` source-fidelity, beating the older full-emblem tile-graph CUDA baseline at `0.0283` @@ -57,17 +57,19 @@ The most provisional areas are: ### Active pass -This pass focused on iteration speed for `tile-graph`, because the current badge loop had become too slow to support quick debugging. +This pass focused on direct control and iteration speed for `tile-graph` and `hybrid`, so the pipeline can run exact user-chosen lattices instead of always paying for the full search-and-rerank path. What landed: -- `pipeline.py` now reuses the chosen low-confidence phase-rerank reconstruction for `tile-graph` and `hybrid` instead of rebuilding the same selected candidate again after selection -- `tests/test_pipeline.py` now has a regression that locks in that reuse behavior for low-confidence `tile-graph` runs +- `pipeline.py` now accepts pinned `target_width` / `target_height` plus optional `phase_x` / `phase_y` so the pipeline can construct a fixed inference result directly +- `inference.py` now exposes `infer_fixed_lattice(...)`, which searches only the requested phase space for one fixed lattice instead of paying the full size search +- `cli.py` and `compare.py` now expose those controls plus `--skip-phase-rerank` +- `tile_graph.py` now has a process-local model cache keyed by source content, lattice choice, device, and build-affecting tile-graph params +- `tests/test_inference.py`, `tests/test_pipeline.py`, `tests/test_tile_graph.py`, and `tests/test_cli.py` now cover fixed-lattice inference, rerank disabling, and cache reuse Current implementation note: -- on the cleaned badge, the current `tile-graph` end-to-end run now lands at about `270.8s` on this machine after reuse, down from about `443.1s` for the earlier stroke-aware badge run -- profiling the selected badge candidate shows that the solver loop is not the problem: the final `tile-graph` reconstruction used about `142.3s`, and about `131.2s` of that sat inside `_extract_source_region_tiles(...)` -- the low-confidence badge rerank is the other major cost center: `_select_phase_candidate(...)` took about `223.9s` on the same fixture, with about `208.1s` of that spent inside reconstruction probes -- `infer_lattice(...)` still takes about `44.1s` on that fixture, so iteration work that does not change the lattice itself should eventually be able to reuse or pin inference too +- the earlier profiling result still stands: on the searched badge path, the solver loop is not the problem; the heavy costs are `infer_lattice(...)`, low-confidence phase rerank probes, and `_extract_source_region_tiles(...)` +- the new direct-control path is specifically for dodging those costs during iteration when we already know which lattice we want to inspect +- on the cleaned badge at pinned `126x126` / phase `(0.0, -0.2)`, a fixed-lattice CUDA `tile-graph` run now took about `10.2s` on the first same-process run and about `2.1s` on the cached rerun, with the same output and a reported `tile_graph_model_cache_hit` on the second pass - the hybrid remains intentionally conservative: it only biases tile-graph's unary cost with the continuous prepass layout and does not yet replace the pairwise objective or the source-region builder - on the cleaned badge, that conservative hybrid still helps: the latest run under `artifacts/badge-hybrid-v2-cuda/` lands at `0.1785`, improving on the current tile-graph badge baseline at `0.1832` - the hybrid is still far behind the stronger continuous badge result (`0.0832`), so the seam looks promising but it is not yet enough to solve the real contour problem by itself @@ -76,7 +78,7 @@ Current implementation note: Next after that: - split tile-graph iteration into a cached model-build phase and a near-free solve phase so weight tuning does not keep paying the `~131s` source-region cutting bill -- add an explicit fast-iteration path that reuses or pins lattice inference when the experiment is only changing tile-graph internals +- extend the current fixed-lattice path so repeated CLI runs can reuse cached model artifacts across processes instead of only within one Python process - port or reformulate `_extract_source_region_tiles(...)` so it no longer spends most of the badge runtime in Python/NumPy window cutting - reduce rerank probe cost by using cached candidate builds where possible or by using a cheaper rerank proxy before full tile-graph reconstruction - decide whether the next hybrid pass should add geometry-aware pairwise terms before changing the source-region builder again diff --git a/src/repixelizer/cli.py b/src/repixelizer/cli.py index 1df03f1..d050b3f 100644 --- a/src/repixelizer/cli.py +++ b/src/repixelizer/cli.py @@ -18,6 +18,10 @@ def build_parser() -> argparse.ArgumentParser: target.add_argument("input", help="Source image path") target.add_argument("--out", required=True, help="Output image path") target.add_argument("--target-size", type=int, default=None, help="Override the inferred target max dimension") + target.add_argument("--target-width", type=int, default=None, help="Pin the output lattice width directly") + target.add_argument("--target-height", type=int, default=None, help="Pin the output lattice height directly") + target.add_argument("--phase-x", type=float, default=None, help="Pin the lattice phase offset on the x axis") + target.add_argument("--phase-y", type=float, default=None, help="Pin the lattice phase offset on the y axis") target.add_argument("--palette", default=None, help="Optional palette file (.gpl, .txt, .json)") target.add_argument( "--palette-mode", @@ -40,6 +44,11 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Remove light neutral edge-connected backgrounds such as fake transparency checkerboards", ) + target.add_argument( + "--skip-phase-rerank", + action="store_true", + help="Skip low-confidence phase reranking and run the selected or pinned lattice directly", + ) run_parser = subparsers.add_parser("run", help="Run the optimizer.") add_shared_arguments(run_parser) @@ -125,6 +134,10 @@ def main(argv: list[str] | None = None) -> int: args.input, args.out, target_size=args.target_size, + target_width=args.target_width, + target_height=args.target_height, + phase_x=args.phase_x, + phase_y=args.phase_y, palette_path=args.palette, palette_mode=args.palette_mode, diagnostics_dir=args.diagnostics_dir, @@ -133,6 +146,7 @@ def main(argv: list[str] | None = None) -> int: device=args.device, reconstruction_mode=args.reconstruction_mode, strip_background=args.strip_background, + enable_phase_rerank=not args.skip_phase_rerank, ) return 0 if command == "benchmark": @@ -173,6 +187,10 @@ def main(argv: list[str] | None = None) -> int: args.input, args.out, target_size=args.target_size, + target_width=args.target_width, + target_height=args.target_height, + phase_x=args.phase_x, + phase_y=args.phase_y, palette_path=args.palette, palette_mode=args.palette_mode, diagnostics_dir=args.diagnostics_dir, @@ -181,5 +199,6 @@ def main(argv: list[str] | None = None) -> int: device=args.device, reconstruction_mode=args.reconstruction_mode, strip_background=args.strip_background, + enable_phase_rerank=not args.skip_phase_rerank, ) return 0 diff --git a/src/repixelizer/compare.py b/src/repixelizer/compare.py index ed3eccb..589d8ce 100644 --- a/src/repixelizer/compare.py +++ b/src/repixelizer/compare.py @@ -57,6 +57,10 @@ def run_compare( output_path: str | Path, *, target_size: int | None = None, + target_width: int | None = None, + target_height: int | None = None, + phase_x: float | None = None, + phase_y: float | None = None, palette_path: str | Path | None = None, palette_mode: str = "off", diagnostics_dir: str | Path | None = None, @@ -65,6 +69,7 @@ def run_compare( device: str = "auto", strip_background: bool = False, reconstruction_mode: str = "continuous", + enable_phase_rerank: bool = True, ) -> dict[str, Any]: diagnostics_path = Path(diagnostics_dir) if diagnostics_dir else Path(output_path).with_suffix("") diagnostics_path.mkdir(parents=True, exist_ok=True) @@ -72,6 +77,10 @@ def run_compare( input_path, output_path, target_size=target_size, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, palette_path=palette_path, palette_mode=palette_mode, diagnostics_dir=diagnostics_path, @@ -80,6 +89,7 @@ def run_compare( device=device, strip_background=strip_background, reconstruction_mode=reconstruction_mode, + enable_phase_rerank=enable_phase_rerank, ) source = result.source_rgba palette = load_palette(palette_path) if palette_path else None diff --git a/src/repixelizer/inference.py b/src/repixelizer/inference.py index 04b5dff..6f03ec3 100644 --- a/src/repixelizer/inference.py +++ b/src/repixelizer/inference.py @@ -410,14 +410,15 @@ def _score_phase_group( prior_cell_x: float, prior_cell_y: float, prior_reliability: float, - phase_values: np.ndarray, + phase_x_values: np.ndarray, + phase_y_values: np.ndarray, device: str, ): torch, F = _require_torch() height, width = rgba.shape[:2] premult_rgba = premultiply(rgba) source_t = torch.from_numpy(premult_rgba.transpose(2, 0, 1)[None, ...]).to(device=device, dtype=torch.float32) - phase_xs, phase_ys = np.meshgrid(phase_values, phase_values, indexing="xy") + phase_xs, phase_ys = np.meshgrid(phase_x_values, phase_y_values, indexing="xy") phase_x_t = torch.from_numpy(phase_xs.reshape(-1).astype(np.float32)).to(device=device) phase_y_t = torch.from_numpy(phase_ys.reshape(-1).astype(np.float32)).to(device=device) batch = phase_x_t.shape[0] @@ -489,7 +490,8 @@ def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str prior_cell_x=prior_cell_x, prior_cell_y=prior_cell_y, prior_reliability=prior_reliability, - phase_values=phase_values, + phase_x_values=phase_values, + phase_y_values=phase_values, device=resolved_device, ) ) @@ -511,6 +513,48 @@ def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str ) +def infer_fixed_lattice( + rgba: np.ndarray, + *, + target_width: int, + target_height: int, + phase_x: float | None = None, + phase_y: float | None = None, + device: str = "auto", +) -> InferenceResult: + torch, _ = _require_torch() + resolved_device = _resolve_device(torch, device) + prior_cell_x, prior_cell_y, prior_reliability = _estimate_lattice_prior_details(rgba) + phase_values = np.linspace(-0.4, 0.4, num=5, dtype=np.float32) + phase_x_values = np.asarray([phase_x], dtype=np.float32) if phase_x is not None else phase_values + phase_y_values = np.asarray([phase_y], dtype=np.float32) if phase_y is not None else phase_values + + candidates = _score_phase_group( + rgba, + target_width=max(1, int(target_width)), + target_height=max(1, int(target_height)), + prior_cell_x=prior_cell_x, + prior_cell_y=prior_cell_y, + prior_reliability=prior_reliability, + phase_x_values=phase_x_values, + phase_y_values=phase_y_values, + device=resolved_device, + ) + candidates.sort(key=lambda item: item.score, reverse=True) + best = candidates[0] + second = candidates[1] if len(candidates) > 1 else best + confidence = max(0.0, best.score - second.score) + top_candidates = candidates[:8] + return InferenceResult( + target_width=best.target_width, + target_height=best.target_height, + phase_x=best.phase_x, + phase_y=best.phase_y, + confidence=confidence, + top_candidates=top_candidates, + ) + + def inference_to_json(result: InferenceResult) -> dict[str, object]: return { "target_width": result.target_width, diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index d86848b..eb3b769 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -17,7 +17,7 @@ from .diagnostics import ( write_lattice_overlay, ) from .discrete import cleanup_pixels -from .inference import infer_lattice, inference_to_json +from .inference import infer_fixed_lattice, infer_lattice, inference_to_json from .io import load_rgba, nearest_resize, save_rgba from .metrics import ( foreground_edge_concentration, @@ -41,6 +41,10 @@ def run_pipeline( output_path: str | Path, *, target_size: int | None = None, + target_width: int | None = None, + target_height: int | None = None, + phase_x: float | None = None, + phase_y: float | None = None, palette_path: str | Path | None = None, palette_mode: str = "off", diagnostics_dir: str | Path | None = None, @@ -50,13 +54,35 @@ def run_pipeline( solver_params: SolverHyperParams | None = None, strip_background: bool = False, reconstruction_mode: str = "continuous", + enable_phase_rerank: bool = True, ) -> RunResult: started = time.perf_counter() solver_params = solver_params or SolverHyperParams() source = load_rgba(input_path) if strip_background: source = strip_edge_background(source) - inference = infer_lattice(source, target_size=target_size, device=device) + fixed_dims = _resolve_requested_target_dims( + source_width=source.shape[1], + source_height=source.shape[0], + target_size=target_size, + target_width=target_width, + target_height=target_height, + phase_x=phase_x, + phase_y=phase_y, + ) + if fixed_dims is None: + inference = infer_lattice(source, target_size=target_size, device=device) + inference_mode = "searched" + else: + inference = infer_fixed_lattice( + source, + target_width=fixed_dims[0], + target_height=fixed_dims[1], + phase_x=phase_x, + phase_y=phase_y, + device=device, + ) + inference_mode = "fixed" analysis = analyze_source(source, seed=seed, device=device if reconstruction_mode in {"tile-graph", "hybrid"} else None) inference, cached_reconstruction = _select_phase_candidate_with_reconstruction( source, @@ -66,6 +92,7 @@ def run_pipeline( device=device, solver_params=solver_params, reconstruction_mode=reconstruction_mode, + enable_phase_rerank=enable_phase_rerank, ) if cached_reconstruction is not None: solver, reconstruction_diagnostics = cached_reconstruction @@ -121,12 +148,18 @@ def run_pipeline( run_json["inference"] = inference_to_json(inference) run_json["settings"] = { "target_size": target_size, + "target_width": target_width, + "target_height": target_height, + "phase_x": phase_x, + "phase_y": phase_y, "palette_mode": palette_mode, "seed": seed, "steps": steps, "device": device, "strip_background": strip_background, "reconstruction_mode": reconstruction_mode, + "enable_phase_rerank": enable_phase_rerank, + "inference_mode": inference_mode, "solver_params": solver_params.to_dict(), } if diagnostics.get("reconstruction"): @@ -146,6 +179,7 @@ def _select_phase_candidate( device: str, solver_params: SolverHyperParams | None = None, reconstruction_mode: str = "continuous", + enable_phase_rerank: bool = True, ) -> InferenceResult: selected, _ = _select_phase_candidate_with_reconstruction( source, @@ -155,6 +189,7 @@ def _select_phase_candidate( device=device, solver_params=solver_params, reconstruction_mode=reconstruction_mode, + enable_phase_rerank=enable_phase_rerank, ) return selected @@ -168,8 +203,11 @@ def _select_phase_candidate_with_reconstruction( device: str, solver_params: SolverHyperParams | None = None, reconstruction_mode: str = "continuous", + enable_phase_rerank: bool = True, ): solver_params = solver_params or SolverHyperParams() + if not enable_phase_rerank: + return inference, None if len(inference.top_candidates) <= 1 or inference.confidence >= solver_params.phase_rerank_confidence_threshold: return inference, None @@ -306,6 +344,41 @@ def _select_phase_candidate_with_reconstruction( return best_candidate, cached +def _resolve_requested_target_dims( + *, + source_width: int, + source_height: int, + target_size: int | None, + target_width: int | None, + target_height: int | None, + phase_x: float | None, + phase_y: float | None, +) -> tuple[int, int] | None: + explicit_size = target_width is not None or target_height is not None + explicit_phase = phase_x is not None or phase_y is not None + if not explicit_size and target_size is None: + if explicit_phase: + raise ValueError("phase_x/phase_y require target_size or target_width/target_height so the lattice can be fixed explicitly.") + return None + if target_size is not None and explicit_size: + raise ValueError("target_size cannot be combined with target_width or target_height.") + if target_size is not None: + if source_width >= source_height: + return int(target_size), max(1, round(source_height * int(target_size) / max(1, source_width))) + return max(1, round(source_width * int(target_size) / max(1, source_height))), int(target_size) + if target_width is None and target_height is None: + return None + if target_width is None: + resolved_height = int(target_height) + resolved_width = max(1, round(source_width * resolved_height / max(1, source_height))) + return resolved_width, resolved_height + if target_height is None: + resolved_width = int(target_width) + resolved_height = max(1, round(source_height * resolved_width / max(1, source_width))) + return resolved_width, resolved_height + return int(target_width), int(target_height) + + def _normalize_penalty(values, *, higher_is_better: bool = False) -> list[float]: raw = np.asarray(list(values), dtype=np.float32) if raw.size == 0: diff --git a/src/repixelizer/tile_graph.py b/src/repixelizer/tile_graph.py index 5cb2131..f3c85b1 100644 --- a/src/repixelizer/tile_graph.py +++ b/src/repixelizer/tile_graph.py @@ -1,7 +1,8 @@ from __future__ import annotations -from collections import deque -from dataclasses import dataclass +import hashlib +from collections import OrderedDict, deque +from dataclasses import dataclass, replace from typing import Any import numpy as np @@ -69,6 +70,70 @@ class TileGraphModel: model_device: str geometry_reference_rgba: np.ndarray | None = None geometry_strength: np.ndarray | None = None + cache_hit: bool = False + + +_TILE_GRAPH_MODEL_CACHE_MAX_ENTRIES = 8 +_TILE_GRAPH_MODEL_CACHE: OrderedDict[tuple[object, ...], TileGraphModel] = OrderedDict() + + +def clear_tile_graph_model_cache() -> None: + _TILE_GRAPH_MODEL_CACHE.clear() + + +def _hash_array(array: np.ndarray) -> str: + contiguous = np.ascontiguousarray(array) + return hashlib.blake2b(contiguous.view(np.uint8), digest_size=16).hexdigest() + + +def _tile_graph_model_cache_key( + source_rgba: np.ndarray, + *, + inference: InferenceResult, + solver_params: SolverHyperParams, + device: str, +) -> tuple[object, ...]: + build_param_names = ( + "alpha_transparent_threshold", + "source_edge_detail_threshold", + "tile_graph_max_candidates_per_coord", + "tile_graph_edge_candidates_per_coord", + "tile_graph_component_color_threshold", + "tile_graph_component_alpha_threshold", + "tile_graph_source_region_stride", + "tile_graph_source_region_min_area_ratio", + "tile_graph_source_region_window_coverage", + "tile_graph_stroke_linearity_threshold", + "tile_graph_stroke_step_scale", + "tile_graph_stroke_minor_limit_scale", + ) + build_signature = tuple((name, getattr(solver_params, name)) for name in build_param_names) + return ( + _hash_array(source_rgba), + source_rgba.shape, + str(source_rgba.dtype), + int(inference.target_width), + int(inference.target_height), + float(inference.phase_x), + float(inference.phase_y), + device, + build_signature, + ) + + +def _get_cached_tile_graph_model(cache_key: tuple[object, ...]) -> TileGraphModel | None: + cached = _TILE_GRAPH_MODEL_CACHE.get(cache_key) + if cached is None: + return None + _TILE_GRAPH_MODEL_CACHE.move_to_end(cache_key) + return cached + + +def _store_cached_tile_graph_model(cache_key: tuple[object, ...], model: TileGraphModel) -> None: + _TILE_GRAPH_MODEL_CACHE[cache_key] = model + _TILE_GRAPH_MODEL_CACHE.move_to_end(cache_key) + while len(_TILE_GRAPH_MODEL_CACHE) > _TILE_GRAPH_MODEL_CACHE_MAX_ENTRIES: + _TILE_GRAPH_MODEL_CACHE.popitem(last=False) def _make_regular_uv(height: int, width: int, target_height: int, target_width: int, phase_x: float, phase_y: float) -> np.ndarray: @@ -752,6 +817,34 @@ def build_tile_graph_model( torch = _require_torch() resolved_device = _resolve_device(torch, device) solver_params = solver_params or SolverHyperParams() + geometry_reference_np = None + if geometry_reference_rgba is not None: + geometry_reference_np = np.asarray(geometry_reference_rgba, dtype=np.float32) + if geometry_reference_np.shape[:2] != (inference.target_height, inference.target_width): + raise ValueError( + "geometry_reference_rgba must match the inferred output size when provided to tile-graph." + ) + geometry_strength_np = None + if geometry_guidance_strength is not None: + geometry_strength_np = np.asarray(geometry_guidance_strength, dtype=np.float32) + if geometry_strength_np.shape != (inference.target_height, inference.target_width): + raise ValueError( + "geometry_guidance_strength must match the inferred output size when provided to tile-graph." + ) + cache_key = _tile_graph_model_cache_key( + source_rgba, + inference=inference, + solver_params=solver_params, + device=resolved_device, + ) + cached_model = _get_cached_tile_graph_model(cache_key) + if cached_model is not None: + return replace( + cached_model, + geometry_reference_rgba=geometry_reference_np, + geometry_strength=geometry_strength_np, + cache_hit=True, + ) source_reference = build_source_lattice_reference( source_rgba, target_width=inference.target_width, @@ -927,22 +1020,7 @@ def build_tile_graph_model( edge_density = float( np.mean(source_reference.edge_strength >= solver_params.source_edge_detail_threshold) ) - geometry_reference_np = None - if geometry_reference_rgba is not None: - geometry_reference_np = np.asarray(geometry_reference_rgba, dtype=np.float32) - if geometry_reference_np.shape[:2] != (inference.target_height, inference.target_width): - raise ValueError( - "geometry_reference_rgba must match the inferred output size when provided to tile-graph." - ) - geometry_strength_np = None - if geometry_guidance_strength is not None: - geometry_strength_np = np.asarray(geometry_guidance_strength, dtype=np.float32) - if geometry_strength_np.shape != (inference.target_height, inference.target_width): - raise ValueError( - "geometry_guidance_strength must match the inferred output size when provided to tile-graph." - ) - - return TileGraphModel( + base_model = TileGraphModel( candidate_rgba=candidate_rgba_np, candidate_coords=candidate_coords_np, candidate_area_ratio=candidate_area_ratio_np, @@ -958,8 +1036,16 @@ def build_tile_graph_model( edge_density=edge_density, average_choices=average_choices, model_device=resolved_device, + geometry_reference_rgba=None, + geometry_strength=None, + cache_hit=False, + ) + _store_cached_tile_graph_model(cache_key, base_model) + return replace( + base_model, geometry_reference_rgba=geometry_reference_np, geometry_strength=geometry_strength_np, + cache_hit=False, ) @@ -1394,6 +1480,7 @@ def optimize_tile_graph( "mode": "tile-graph", "tile_graph_model_device": model.model_device, "tile_graph_solver_device": resolved_device, + "tile_graph_model_cache_hit": model.cache_hit, "tile_graph_proposal_mode": "atomic-components", "tile_graph_component_count": model.component_count, "tile_graph_candidate_count": int(model.candidate_rgba.shape[0]), diff --git a/tests/test_cli.py b/tests/test_cli.py index 9539062..ae585a9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -27,3 +27,31 @@ def test_default_command_accepts_strip_background_flag(tmp_path: Path) -> None: exit_code = main([str(input_path), "--out", str(output_path), "--steps", "0", "--strip-background"]) assert exit_code == 0 assert output_path.exists() + + +def test_default_command_accepts_fixed_lattice_flags(tmp_path: Path) -> None: + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=8, phase_x=0.2, phase_y=0.25, blur_radius=0.4) + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + save_rgba(input_path, fake) + exit_code = main( + [ + str(input_path), + "--out", + str(output_path), + "--steps", + "0", + "--target-width", + "16", + "--target-height", + "16", + "--phase-x", + "0.2", + "--phase-y", + "0.25", + "--skip-phase-rerank", + ] + ) + assert exit_code == 0 + assert output_path.exists() diff --git a/tests/test_inference.py b/tests/test_inference.py index ab2bf55..6934605 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -6,6 +6,7 @@ from repixelizer.inference import ( _combine_axis_priors, _hint_target_sizes_from_spacing, _top_candidates_by_size, + infer_fixed_lattice, infer_lattice, ) from repixelizer.types import InferenceCandidate @@ -67,3 +68,37 @@ def test_infer_lattice_recovers_emblem_scale() -> None: assert result.target_width in range(28, 37) assert result.target_height in range(28, 37) assert result.confidence >= 0.0 + + +def test_infer_fixed_lattice_honors_exact_size_and_phase() -> None: + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=10, phase_x=0.2, phase_y=-0.2, blur_radius=0.4, seed=9) + result = infer_fixed_lattice( + fake, + target_width=16, + target_height=16, + phase_x=0.2, + phase_y=-0.2, + device="cpu", + ) + assert result.target_width == 16 + assert result.target_height == 16 + assert abs(result.phase_x - 0.2) <= 1e-6 + assert abs(result.phase_y + 0.2) <= 1e-6 + assert len(result.top_candidates) == 1 + + +def test_infer_fixed_lattice_searches_phase_within_pinned_size() -> None: + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=8, phase_x=0.2, phase_y=0.25, blur_radius=0.45, seed=3) + result = infer_fixed_lattice( + fake, + target_width=16, + target_height=16, + device="cpu", + ) + assert result.target_width == 16 + assert result.target_height == 16 + assert len(result.top_candidates) > 1 + assert all(candidate.target_width == 16 for candidate in result.top_candidates) + assert all(candidate.target_height == 16 for candidate in result.top_candidates) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index ba46eef..485ee07 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -322,6 +322,35 @@ def test_pipeline_hybrid_mode_writes_reconstruction_diagnostics(tmp_path: Path) assert result.diagnostics["reconstruction"]["mode"] == "hybrid" +def test_select_phase_candidate_can_skip_phase_rerank(monkeypatch) -> None: + source = np.zeros((8, 8, 4), dtype=np.float32) + candidate_a = InferenceCandidate(target_width=8, target_height=8, phase_x=0.0, phase_y=0.0, score=0.9, breakdown={}) + candidate_b = InferenceCandidate(target_width=10, target_height=10, phase_x=0.2, phase_y=0.2, score=0.89, breakdown={}) + inference = InferenceResult( + target_width=8, + target_height=8, + phase_x=0.0, + phase_y=0.0, + confidence=0.0, + top_candidates=[candidate_a, candidate_b], + ) + + def fail(*args, **kwargs): + raise AssertionError("phase rerank probe should not run when disabled") + + monkeypatch.setattr("repixelizer.pipeline.optimize_uv_field", fail) + selected = _select_phase_candidate( + source, + inference, + analysis=object(), + seed=7, + device="cpu", + enable_phase_rerank=False, + ) + assert selected.target_width == candidate_a.target_width + assert selected.target_height == candidate_a.target_height + + def test_pipeline_reuses_phase_probe_reconstruction_for_tile_graph(tmp_path: Path, monkeypatch) -> None: source = np.zeros((6, 6, 4), dtype=np.float32) source[..., 3] = 1.0 @@ -401,6 +430,69 @@ def test_pipeline_reuses_phase_probe_reconstruction_for_tile_graph(tmp_path: Pat assert result.output_rgba.shape[:2] == (3, 3) +def test_pipeline_fixed_target_and_phase_skip_search_inference(tmp_path: Path, monkeypatch) -> None: + source = np.zeros((6, 6, 4), dtype=np.float32) + source[..., 3] = 1.0 + input_path = tmp_path / "input.png" + output_path = tmp_path / "output.png" + + from repixelizer.io import save_rgba + + save_rgba(input_path, source) + + inference = InferenceResult( + target_width=3, + target_height=3, + phase_x=0.2, + phase_y=-0.2, + confidence=1.0, + top_candidates=[], + ) + + class DummyAnalysis: + cluster_preview = np.zeros((6, 6, 4), dtype=np.float32) + + def fail_infer(*args, **kwargs): + raise AssertionError("full lattice search should not run for fixed target/phase") + + def fake_infer_fixed(*args, **kwargs): + return inference + + def fake_optimize_uv_field(source_rgba, inference, analysis, steps, seed, device, solver_params=None): + rgba = np.zeros((inference.target_height, inference.target_width, 4), dtype=np.float32) + rgba[..., 3] = 1.0 + return SolverArtifacts( + target_rgba=rgba, + uv_field=np.zeros((inference.target_height, inference.target_width, 2), dtype=np.float32), + guidance_strength=np.zeros((inference.target_height, inference.target_width), dtype=np.float32), + initial_rgba=rgba.copy(), + loss_history=[], + ) + + monkeypatch.setattr("repixelizer.pipeline.infer_lattice", fail_infer) + monkeypatch.setattr("repixelizer.pipeline.infer_fixed_lattice", fake_infer_fixed) + monkeypatch.setattr("repixelizer.pipeline.analyze_source", lambda *args, **kwargs: DummyAnalysis()) + monkeypatch.setattr("repixelizer.pipeline.optimize_uv_field", fake_optimize_uv_field) + + result = run_pipeline( + input_path, + output_path, + target_width=3, + target_height=3, + phase_x=0.2, + phase_y=-0.2, + enable_phase_rerank=False, + steps=0, + device="cpu", + ) + + assert output_path.exists() + assert result.inference.target_width == 3 + assert result.inference.target_height == 3 + assert result.inference.phase_x == 0.2 + assert result.inference.phase_y == -0.2 + + def test_phase_rerank_can_accept_low_confidence_size_jump_with_strong_support(monkeypatch) -> None: source = np.zeros((16, 16, 4), dtype=np.float32) candidate_a = InferenceCandidate(target_width=16, target_height=16, phase_x=0.0, phase_y=0.0, score=0.91, breakdown={}) diff --git a/tests/test_tile_graph.py b/tests/test_tile_graph.py index 98b4cf5..c478c21 100644 --- a/tests/test_tile_graph.py +++ b/tests/test_tile_graph.py @@ -12,7 +12,7 @@ from repixelizer.metrics import source_lattice_consistency_breakdown from repixelizer.params import SolverHyperParams from repixelizer.pipeline import run_pipeline from repixelizer.synthetic import fake_pixelize, make_emblem -from repixelizer.tile_graph import build_tile_graph_model, optimize_tile_graph +from repixelizer.tile_graph import build_tile_graph_model, clear_tile_graph_model_cache, optimize_tile_graph from repixelizer.types import InferenceResult @@ -361,6 +361,38 @@ def test_tile_graph_cuda_matches_cpu_on_small_case() -> None: assert np.allclose(cpu_artifacts.target_rgba, cuda_artifacts.target_rgba, atol=1e-5) +def test_tile_graph_model_cache_reuses_fixed_lattice_build() -> None: + clear_tile_graph_model_cache() + source = make_emblem(16, 16) + fake = fake_pixelize(source, upscale=8, phase_x=0.1, phase_y=-0.05, blur_radius=0.4, warp_strength=0.15, warp_detail=4) + inference = InferenceResult( + target_width=16, + target_height=16, + phase_x=0.1, + phase_y=-0.05, + confidence=1.0, + top_candidates=[], + ) + analysis = analyze_source(fake, seed=7, device="cpu") + + first = build_tile_graph_model( + fake, + inference=inference, + analysis=analysis, + device="cpu", + ) + second = build_tile_graph_model( + fake, + inference=inference, + analysis=analysis, + device="cpu", + ) + + assert first.cache_hit is False + assert second.cache_hit is True + assert np.array_equal(first.candidate_rgba, second.candidate_rgba) + + def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Path) -> None: source = make_emblem(24, 24) fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.2, blur_radius=0.6, warp_strength=0.2, warp_detail=5) @@ -385,6 +417,7 @@ def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Pa assert run_json["settings"]["reconstruction_mode"] == "tile-graph" assert run_json["reconstruction"]["mode"] == "tile-graph" assert run_json["reconstruction"]["tile_graph_model_device"] == "cpu" + assert "tile_graph_model_cache_hit" in run_json["reconstruction"] assert run_json["reconstruction"]["tile_graph_candidate_count"] > 0 assert set(run_json["source_fidelity"].keys()) == {"snap_initial", "solver_target", "final_output"} assert result.diagnostics["reconstruction"]["mode"] == "tile-graph"