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

Trim solver syncs and sparse GUI previews

Browse source
This commit is contained in:
Meta Vens 2026-04-25 05:33:28 +01:00
commit 67ddcceff0
4 changed files with 152 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/repixelizer/gui.py
View file

@ -136,6 +136,8 @@ class GuiJob:
job_id: str
filename: str
created_at: float = field(default_factory=time.time)
phase_field_preview_stride: int = 4
phase_field_include_snapshot: bool = True
status: str = "queued"
error: str | None = None
events: list[dict[str, Any]] = field(default_factory=list)

164
src/repixelizer/phase_field.py
View file

@ -330,11 +330,11 @@ def _phase_field_loss(torch, F, prep: _PhaseFieldPrep, disp_t, solver_params: So
+ magnitude * solver_params.phase_field_magnitude_weight
)
terms = {
"local_coherence": float(local_coherence.detach().cpu().item()),
"local_edge": float(local_edge.detach().cpu().item()),
"smoothness": float(smoothness.detach().cpu().item()),
"collapse": float(collapse.detach().cpu().item()),
"magnitude": float(magnitude.detach().cpu().item()),
"local_coherence": local_coherence.detach(),
"local_edge": local_edge.detach(),
"smoothness": smoothness.detach(),
"collapse": collapse.detach(),
"magnitude": magnitude.detach(),
}
return loss, sampled_rgba, terms
@ -368,6 +368,65 @@ def _observer_snapshot(
}
def _materialize_phase_terms(terms: dict[str, object]) -> dict[str, float]:
materialized: dict[str, float] = {}
for key, value in terms.items():
if hasattr(value, "detach"):
materialized[key] = float(value.detach().cpu().item())
else:
materialized[key] = float(value)
return materialized
def _materialize_loss_history(torch, values: list[object]) -> list[float]:
if not values:
return []
normalized = []
for value in values:
if hasattr(value, "detach"):
normalized.append(value.detach().reshape(()))
else:
normalized.append(torch.tensor(float(value)))
stacked = torch.stack(normalized)
return [float(item) for item in stacked.cpu().tolist()]
def _observer_option(observer: PipelineObserver | None, name: str, default: object) -> object:
if observer is None:
return default
direct = getattr(observer, name, None)
if direct is not None:
return direct
owner = getattr(observer, "__self__", None)
if owner is not None:
owner_value = getattr(owner, name, None)
if owner_value is not None:
return owner_value
return default
def _observer_preview_stride(observer: PipelineObserver | None) -> int:
raw = _observer_option(observer, "phase_field_preview_stride", 1)
try:
return max(1, int(raw))
except (TypeError, ValueError):
return 1
def _observer_needs_phase_field_snapshot(observer: PipelineObserver | None) -> bool:
return bool(_observer_option(observer, "phase_field_include_snapshot", True))
def _should_emit_phase_field_step(step: int, total_steps: int, *, preview_stride: int) -> bool:
if total_steps <= 0:
return True
if preview_stride <= 1:
return True
if step >= total_steps:
return False
return step % max(1, preview_stride) == 0
def optimize_phase_field(
rgba: np.ndarray,
inference: InferenceResult,
@ -384,6 +443,8 @@ def optimize_phase_field(
torch.manual_seed(seed)
if device == "cuda":
torch.cuda.manual_seed_all(seed)
preview_stride = _observer_preview_stride(observer)
include_snapshots = _observer_needs_phase_field_snapshot(observer)
prep = _prepare_phase_field(
torch,
@ -394,16 +455,16 @@ def optimize_phase_field(
solver_params,
device=device,
)
emit_observer(
observer,
"phase_field_prepared",
uv0_px=prep.uv0_px_t[0].detach().cpu().numpy().astype(np.float32),
guidance=prep.guidance.copy(),
cell_x=float(prep.cell_x),
cell_y=float(prep.cell_y),
target_width=int(inference.target_width),
target_height=int(inference.target_height),
)
prepared_payload = {
"cell_x": float(prep.cell_x),
"cell_y": float(prep.cell_y),
"target_width": int(inference.target_width),
"target_height": int(inference.target_height),
}
if include_snapshots:
prepared_payload["uv0_px"] = prep.uv0_px_t[0].detach().cpu().numpy().astype(np.float32)
prepared_payload["guidance"] = prep.guidance.copy()
emit_observer(observer, "phase_field_prepared", **prepared_payload)
disp_t = torch.zeros_like(prep.uv0_px_t, device=device, dtype=torch.float32, requires_grad=True)
optimizer = torch.optim.Adam([disp_t], lr=solver_params.phase_field_learning_rate)
@ -411,16 +472,16 @@ def optimize_phase_field(
initial_x = np.rint(prep.uv0_px_t[0, ..., 0].detach().cpu().numpy()).astype(np.int32).clip(0, prep.width - 1)
initial_y = np.rint(prep.uv0_px_t[0, ..., 1].detach().cpu().numpy()).astype(np.int32).clip(0, prep.height - 1)
initial_rgba = _nearest_source_rgba(rgba, initial_x, initial_y)
emit_observer(
observer,
"phase_field_initial",
step=0,
total_steps=int(max(0, steps)),
**_observer_snapshot(prep, rgba, prep.uv0_px_t[0].detach().cpu().numpy().astype(np.float32)),
)
initial_payload = {
"step": 0,
"total_steps": int(max(0, steps)),
}
if include_snapshots:
initial_payload.update(_observer_snapshot(prep, rgba, prep.uv0_px_t[0].detach().cpu().numpy().astype(np.float32)))
emit_observer(observer, "phase_field_initial", **initial_payload)
loss_history: list[float] = []
final_terms: dict[str, float] = {}
loss_history_raw: list[object] = []
final_terms_raw: dict[str, object] = {}
min_dx = solver_params.phase_field_min_spacing_ratio * prep.cell_x
min_dy = solver_params.phase_field_min_spacing_ratio * prep.cell_y
max_dx = solver_params.phase_field_max_displacement_ratio * prep.cell_x
@ -443,30 +504,33 @@ def optimize_phase_field(
width=prep.width,
height=prep.height,
)
loss_value = float(loss.detach().cpu().item())
loss_history.append(loss_value)
final_terms = terms
if observer is not None:
current_pos = (prep.uv0_px_t + disp_t)[0].detach().cpu().numpy().astype(np.float32)
emit_observer(
observer,
"phase_field_step",
step=int(step_index + 1),
total_steps=int(max(0, steps)),
loss=loss_value,
terms=terms.copy(),
**_observer_snapshot(prep, rgba, current_pos),
)
loss_detached = loss.detach()
loss_history_raw.append(loss_detached)
final_terms_raw = terms
step_number = int(step_index + 1)
if observer is not None and _should_emit_phase_field_step(step_number, int(max(0, steps)), preview_stride=preview_stride):
payload = {
"step": step_number,
"total_steps": int(max(0, steps)),
"loss": float(loss_detached.cpu().item()),
"terms": _materialize_phase_terms(terms),
}
if include_snapshots:
current_pos = (prep.uv0_px_t + disp_t)[0].detach().cpu().numpy().astype(np.float32)
payload.update(_observer_snapshot(prep, rgba, current_pos))
emit_observer(observer, "phase_field_step", **payload)
with torch.no_grad():
if steps <= 0:
loss, _sampled_rgba, final_terms = _phase_field_loss(torch, F, prep, disp_t, solver_params)
loss_history.append(float(loss.detach().cpu().item()))
loss, _sampled_rgba, final_terms_raw = _phase_field_loss(torch, F, prep, disp_t, solver_params)
loss_history_raw.append(loss.detach())
final_px = prep.uv0_px_t + disp_t
final_x = torch.round(final_px[..., 0]).clamp(0, prep.width - 1).to(dtype=torch.long)
final_y = torch.round(final_px[..., 1]).clamp(0, prep.height - 1).to(dtype=torch.long)
final_pos_norm = _pixel_to_normalized(final_px.clone(), width=prep.width, height=prep.height)[0].detach().cpu().numpy()
loss_history = _materialize_loss_history(torch, loss_history_raw)
final_terms = _materialize_phase_terms(final_terms_raw)
final_x_np = final_x[0].detach().cpu().numpy().astype(np.int32)
final_y_np = final_y[0].detach().cpu().numpy().astype(np.int32)
target_rgba = _nearest_source_rgba(rgba, final_x_np, final_y_np)
@ -495,15 +559,17 @@ def optimize_phase_field(
**final_terms,
},
}
emit_observer(
observer,
"phase_field_final",
step=int(max(0, steps)),
total_steps=int(max(0, steps)),
phase_metrics=stage_diagnostics["phase_field"].copy(),
loss_history=loss_history.copy(),
**_observer_snapshot(prep, rgba, final_px[0].detach().cpu().numpy().astype(np.float32)),
)
final_payload = {
"step": int(max(0, steps)),
"total_steps": int(max(0, steps)),
"phase_metrics": stage_diagnostics["phase_field"].copy(),
"loss_history": loss_history.copy(),
}
if loss_history:
final_payload["loss"] = float(loss_history[-1])
if include_snapshots:
final_payload.update(_observer_snapshot(prep, rgba, final_px[0].detach().cpu().numpy().astype(np.float32)))
emit_observer(observer, "phase_field_final", **final_payload)
return SolverArtifacts(
target_rgba=target_rgba,
uv_field=final_pos_norm,

2
src/repixelizer/pipeline.py
View file

@ -417,6 +417,8 @@ def _select_phase_candidate_with_reconstruction(
**common,
)
rerank_observer.phase_field_include_snapshot = False # type: ignore[attr-defined]
candidate_artifacts, _candidate_diagnostics = _run_reconstruction(
source,
inference=candidate_inference,

33
tests/test_gui.py
View file

@ -83,6 +83,39 @@ def test_run_pipeline_rgba_emits_observer_events_for_gui() -> None:
assert events[-4:] == ["cleanup_completed", "stage_started", "palette_completed", "pipeline_completed"]
def test_run_pipeline_rgba_respects_phase_field_preview_stride() -> None:
source = make_emblem(12, 12)
fake = fake_pixelize(source, upscale=6, phase_x=0.12, phase_y=-0.08, blur_radius=0.35, seed=5)
class PreviewObserver:
phase_field_preview_stride = 2
def __init__(self) -> None:
self.events: list[tuple[str, dict[str, object]]] = []
def __call__(self, event: str, payload: dict[str, object]) -> None:
self.events.append((event, payload))
observer = PreviewObserver()
run_pipeline_rgba(
fake,
target_width=12,
target_height=12,
phase_x=0.0,
phase_y=0.0,
steps=5,
device="cpu",
enable_phase_rerank=False,
observer=observer,
)
phase_steps = [payload["step"] for event, payload in observer.events if event == "phase_field_step"]
assert phase_steps == [2, 4]
final_payload = next(payload for event, payload in observer.events if event == "phase_field_final")
assert final_payload["step"] == 5
def test_repo_gui_runner_dispatches_to_gui_main(monkeypatch, capsys) -> None:
repo_root = Path(__file__).resolve().parents[1]
script_path = repo_root / "scripts" / "run_gui.py"