Prune tile-graph sidecar paths

This commit is contained in:
Meta Vens 2026-04-23 03:46:13 +01:00
commit d33a68c310
8 changed files with 83 additions and 291 deletions

View file

@ -36,7 +36,7 @@ def build_parser() -> argparse.ArgumentParser:
target.add_argument(
"--reconstruction-mode",
default="continuous",
choices=("continuous", "tile-graph", "hybrid"),
choices=("continuous", "tile-graph"),
help="Reconstruction engine to run after lattice estimation",
)
target.add_argument(

View file

@ -92,8 +92,6 @@ class SolverHyperParams:
tile_graph_coverage_weight: float = 0.05
tile_graph_delta_weight: float = 0.45
tile_graph_iterations: int = 10
hybrid_geometry_match_weight: float = 0.35
hybrid_geometry_edge_boost: float = 0.55
def to_dict(self) -> dict[str, float | int]:
return asdict(self)

View file

@ -32,10 +32,6 @@ from .tile_graph import optimize_tile_graph
from .types import InferenceResult, RunResult
def _reuse_phase_probe_reconstruction(reconstruction_mode: str) -> bool:
return reconstruction_mode in {"tile-graph", "hybrid"}
def run_pipeline(
input_path: str | Path,
output_path: str | Path,
@ -83,8 +79,8 @@ def run_pipeline(
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(
analysis = analyze_source(source, seed=seed, device=device if reconstruction_mode == "tile-graph" else None)
inference = _select_phase_candidate(
source,
inference,
analysis=analysis,
@ -94,21 +90,16 @@ def run_pipeline(
reconstruction_mode=reconstruction_mode,
enable_phase_rerank=enable_phase_rerank,
)
if cached_reconstruction is not None:
solver, reconstruction_diagnostics = cached_reconstruction
reconstruction_diagnostics = dict(reconstruction_diagnostics)
reconstruction_diagnostics["reused_phase_probe_reconstruction"] = True
else:
solver, reconstruction_diagnostics = _run_reconstruction(
source,
inference=inference,
analysis=analysis,
steps=steps,
seed=seed,
device=device,
solver_params=solver_params,
reconstruction_mode=reconstruction_mode,
)
solver, reconstruction_diagnostics = _run_reconstruction(
source,
inference=inference,
analysis=analysis,
steps=steps,
seed=seed,
device=device,
solver_params=solver_params,
reconstruction_mode=reconstruction_mode,
)
cleanup = cleanup_pixels(solver.target_rgba, source_guidance=solver.guidance_strength)
palette = load_palette(palette_path) if palette_path else None
palette_result = quantize_rgba(cleanup.cleaned_rgba, mode=palette_mode, palette=palette)
@ -181,7 +172,7 @@ def _select_phase_candidate(
reconstruction_mode: str = "continuous",
enable_phase_rerank: bool = True,
) -> InferenceResult:
selected, _ = _select_phase_candidate_with_reconstruction(
return _select_phase_candidate_with_reconstruction(
source,
inference,
analysis=analysis,
@ -191,7 +182,6 @@ def _select_phase_candidate(
reconstruction_mode=reconstruction_mode,
enable_phase_rerank=enable_phase_rerank,
)
return selected
def _select_phase_candidate_with_reconstruction(
@ -204,12 +194,14 @@ def _select_phase_candidate_with_reconstruction(
solver_params: SolverHyperParams | None = None,
reconstruction_mode: str = "continuous",
enable_phase_rerank: bool = True,
):
) -> InferenceResult:
solver_params = solver_params or SolverHyperParams()
if not enable_phase_rerank:
return inference, None
return inference
if reconstruction_mode != "continuous":
return inference
if len(inference.top_candidates) <= 1 or inference.confidence >= solver_params.phase_rerank_confidence_threshold:
return inference, None
return inference
top_score = float(inference.top_candidates[0].score)
candidate_records: list[dict[str, float | InferenceResult]] = []
@ -222,7 +214,7 @@ def _select_phase_candidate_with_reconstruction(
confidence=inference.confidence,
top_candidates=inference.top_candidates,
)
candidate_artifacts, candidate_diagnostics = _run_reconstruction(
candidate_artifacts, _candidate_diagnostics = _run_reconstruction(
source,
inference=candidate_inference,
analysis=analysis,
@ -250,12 +242,11 @@ def _select_phase_candidate_with_reconstruction(
"edge_concentration": foreground_edge_concentration(candidate_artifacts.target_rgba),
"inference_penalty": top_score - float(candidate.score),
"size_delta_ratio": _size_delta_ratio(inference, candidate_inference),
"cached_reconstruction": (candidate_artifacts, candidate_diagnostics),
}
)
if not candidate_records:
return inference, None
return inference
support_penalty = _normalize_penalty(record["support_score"] for record in candidate_records)
edge_position_penalty = _normalize_penalty(record["edge_position_error"] for record in candidate_records)
@ -273,8 +264,6 @@ def _select_phase_candidate_with_reconstruction(
baseline_rank: float | None = None
best_rank = float("inf")
best_candidate = inference
best_cached_reconstruction = None
baseline_cached_reconstruction = None
annotated_candidates = []
for index, record in enumerate(candidate_records):
candidate_inference = record["inference"]
@ -310,11 +299,9 @@ def _select_phase_candidate_with_reconstruction(
)
if baseline_rank is None:
baseline_rank = rank
baseline_cached_reconstruction = record.get("cached_reconstruction")
if rank < best_rank:
best_rank = rank
best_candidate = record["inference"]
best_cached_reconstruction = record.get("cached_reconstruction")
annotated_candidates.sort(key=lambda candidate: float(candidate.breakdown.get("phase_rerank_score", float("inf"))))
for rank, candidate in enumerate(annotated_candidates, start=1):
candidate.breakdown["phase_rerank_rank"] = float(rank)
@ -337,11 +324,9 @@ def _select_phase_candidate_with_reconstruction(
confidence=inference.confidence,
top_candidates=annotated_candidates,
)
cached = baseline_cached_reconstruction if _reuse_phase_probe_reconstruction(reconstruction_mode) else None
return selected, cached
return inference, None
cached = best_cached_reconstruction if _reuse_phase_probe_reconstruction(reconstruction_mode) else None
return best_candidate, cached
return selected
return inference
return best_candidate
def _resolve_requested_target_dims(
@ -420,43 +405,6 @@ def _run_reconstruction(
device=device,
solver_params=solver_params,
)
if reconstruction_mode == "hybrid":
geometry_solver = optimize_uv_field(
source,
inference=inference,
analysis=analysis,
steps=0,
seed=seed,
device=device,
solver_params=solver_params,
)
hybrid_solver, hybrid_diagnostics = optimize_tile_graph(
source,
inference=inference,
analysis=analysis,
steps=steps,
seed=seed,
device=device,
solver_params=solver_params,
geometry_reference_rgba=geometry_solver.target_rgba,
geometry_guidance_strength=geometry_solver.guidance_strength,
)
geometry_score = source_lattice_consistency_breakdown(
source,
geometry_solver.target_rgba,
target_width=inference.target_width,
target_height=inference.target_height,
phase_x=inference.phase_x,
phase_y=inference.phase_y,
)["score"]
diagnostics = dict(hybrid_diagnostics)
diagnostics["mode"] = "hybrid"
diagnostics["hybrid_geometry_prepass_mode"] = "continuous"
diagnostics["hybrid_geometry_prepass_steps"] = 0.0
diagnostics["hybrid_geometry_prepass_source_fidelity"] = float(geometry_score)
diagnostics["hybrid_geometry_match_weight"] = float(solver_params.hybrid_geometry_match_weight)
diagnostics["hybrid_geometry_edge_boost"] = float(solver_params.hybrid_geometry_edge_boost)
return hybrid_solver, diagnostics
solver = optimize_uv_field(
source,
inference=inference,

View file

@ -68,8 +68,6 @@ class TileGraphModel:
edge_density: float
average_choices: float
model_device: str
geometry_reference_rgba: np.ndarray | None = None
geometry_strength: np.ndarray | None = None
cache_hit: bool = False
@ -877,26 +875,10 @@ def build_tile_graph_model(
analysis: SourceAnalysis,
solver_params: SolverHyperParams | None = None,
device: str = "cpu",
geometry_reference_rgba: np.ndarray | None = None,
geometry_guidance_strength: np.ndarray | None = None,
) -> TileGraphModel:
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,
@ -905,12 +887,7 @@ def build_tile_graph_model(
)
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,
)
return replace(cached_model, cache_hit=True)
source_reference = build_source_lattice_reference(
source_rgba,
target_width=inference.target_width,
@ -1106,17 +1083,10 @@ 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,
)
return replace(base_model, cache_hit=False)
def _cell_candidate_span(model: TileGraphModel, y: int, x: int) -> tuple[int, int]:
@ -1175,18 +1145,6 @@ def _tile_graph_unary_cost(model: TileGraphModel, solver_params: SolverHyperPara
+ alpha_error * solver_params.tile_graph_alpha_weight
+ coverage_error * solver_params.tile_graph_coverage_weight
)
if model.geometry_reference_rgba is not None:
geometry_reference = model.geometry_reference_rgba[y, x]
geometry_error = np.mean(np.abs(geometry_reference[None, :] - candidates), axis=-1)
geometry_strength = 1.0
if model.geometry_strength is not None:
geometry_strength += (
float(np.clip(model.geometry_strength[y, x], 0.0, 1.0))
* solver_params.hybrid_geometry_edge_boost
)
total_cost = total_cost + geometry_error * (
solver_params.hybrid_geometry_match_weight * geometry_strength
)
unary_cost[start:end] = total_cost.astype(np.float32)
return unary_cost
@ -1208,17 +1166,6 @@ def _tile_graph_unary_cost_torch(
reference_mean_t = torch.from_numpy(model.reference_mean_rgba).to(device=device, dtype=torch.float32)
reference_edge_t = torch.from_numpy(model.reference_edge_rgba).to(device=device, dtype=torch.float32)
edge_strength_t = torch.from_numpy(model.edge_strength).to(device=device, dtype=torch.float32)
geometry_reference_t = (
torch.from_numpy(model.geometry_reference_rgba).to(device=device, dtype=torch.float32)
if model.geometry_reference_rgba is not None
else None
)
geometry_strength_t = (
torch.from_numpy(model.geometry_strength).to(device=device, dtype=torch.float32)
if model.geometry_strength is not None
else None
)
choice_rgba_t = candidate_rgba_t[choice_indices_t]
choice_deltas_t = candidate_deltas_t[choice_indices_t]
sharp_error_t = (reference_sharp_t[..., None, :] - choice_rgba_t).abs().mean(dim=-1)
@ -1240,14 +1187,6 @@ def _tile_graph_unary_cost_torch(
+ alpha_error * solver_params.tile_graph_alpha_weight
+ coverage_error * solver_params.tile_graph_coverage_weight
)
if geometry_reference_t is not None:
geometry_error_t = (geometry_reference_t[..., None, :] - choice_rgba_t).abs().mean(dim=-1)
geometry_scale_t = 1.0
if geometry_strength_t is not None:
geometry_scale_t = 1.0 + geometry_strength_t[..., None].clamp(0.0, 1.0) * solver_params.hybrid_geometry_edge_boost
unary_cost_t = unary_cost_t + geometry_error_t * (
solver_params.hybrid_geometry_match_weight * geometry_scale_t
)
unary_cost_t = unary_cost_t.masked_fill(~choice_mask_t, float("inf"))
return unary_cost_t, candidate_rgba_t, candidate_deltas_t, choice_rgba_t, choice_deltas_t
@ -1420,8 +1359,6 @@ def optimize_tile_graph(
seed: int,
device: str,
solver_params: SolverHyperParams | None = None,
geometry_reference_rgba: np.ndarray | None = None,
geometry_guidance_strength: np.ndarray | None = None,
) -> tuple[SolverArtifacts, dict[str, Any]]:
del steps
torch = _require_torch()
@ -1436,8 +1373,6 @@ def optimize_tile_graph(
analysis=analysis,
solver_params=solver_params,
device=resolved_device,
geometry_reference_rgba=geometry_reference_rgba,
geometry_guidance_strength=geometry_guidance_strength,
)
choice_indices_np, choice_mask_np = _build_choice_grid(model)
choice_indices_t = torch.from_numpy(choice_indices_np).to(device=resolved_device, dtype=torch.long)
@ -1561,7 +1496,6 @@ def optimize_tile_graph(
"tile_graph_initial_source_fidelity": float(initial_source_fidelity),
"tile_graph_final_source_fidelity": float(final_source_fidelity),
"tile_graph_kept_initial_assignment": kept_initial_assignment,
"tile_graph_geometry_prior": model.geometry_reference_rgba is not None,
}
return (
SolverArtifacts(