GameCult/repixelizer
Watch
9
0
Fork
You've already forked repixelizer
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add fixed lattice tile-graph controls

Browse source
This commit is contained in:
Meta Vens 2026-04-23 02:00:21 +01:00
commit 185cbdf714
11 changed files with 459 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

3
README.md
View file

@ -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:

20
docs/implementation-plan.md
View file

@ -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

19
src/repixelizer/cli.py
View file

@ -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

10
src/repixelizer/compare.py
View file

@ -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

50
src/repixelizer/inference.py
View file

@ -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,

77
src/repixelizer/pipeline.py
View file

@ -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:

123
src/repixelizer/tile_graph.py
View file

@ -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]),

28
tests/test_cli.py
View file

@ -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()

35
tests/test_inference.py
View file

@ -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)

92
tests/test_pipeline.py
View file

@ -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={})

35
tests/test_tile_graph.py
View file

@ -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"