Report lattice search and rerank progress
This commit is contained in:
Meta Vens
parent
99c4f54f7a
commit
a2cf44b3ba
6 changed files with 579 additions and 23 deletions
|
|
@ -208,6 +208,55 @@ class GuiJob:
|
|||
"candidateCount": int(payload["candidate_count"]),
|
||||
"confidence": float(payload["confidence"]),
|
||||
}
|
||||
if event == "lattice_search_started":
|
||||
return {
|
||||
"candidateCount": int(payload["candidate_count"]),
|
||||
"phaseSampleCount": int(payload["phase_sample_count"]),
|
||||
"device": str(payload["device"]),
|
||||
}
|
||||
if event == "lattice_search_progress":
|
||||
return {
|
||||
"completedCandidates": int(payload["completed_candidates"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseSampleCount": int(payload["phase_sample_count"]),
|
||||
"bestScore": None if payload.get("best_score") is None else float(payload["best_score"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_started":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"previewSteps": int(payload["preview_steps"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_step":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"step": int(payload["step"]),
|
||||
"totalSteps": int(payload["total_steps"]),
|
||||
"loss": None if payload.get("loss") is None else float(payload["loss"]),
|
||||
}
|
||||
if event == "phase_rerank_candidate_completed":
|
||||
return {
|
||||
"candidateIndex": int(payload["candidate_index"]),
|
||||
"totalCandidates": int(payload["total_candidates"]),
|
||||
"completedCandidates": int(payload["completed_candidates"]),
|
||||
"targetWidth": int(payload["target_width"]),
|
||||
"targetHeight": int(payload["target_height"]),
|
||||
"phaseX": float(payload["phase_x"]),
|
||||
"phaseY": float(payload["phase_y"]),
|
||||
"totalSteps": int(payload["total_steps"]),
|
||||
"finalLoss": None if payload.get("final_loss") is None else float(payload["final_loss"]),
|
||||
}
|
||||
if event == "stage_started":
|
||||
return {
|
||||
"stage": str(payload["stage"]),
|
||||
|
|
|
|||
|
|
@ -61,8 +61,13 @@ const eventTypes = [
|
|||
"stage_started",
|
||||
"source_loaded",
|
||||
"preprocess_completed",
|
||||
"lattice_search_started",
|
||||
"lattice_search_progress",
|
||||
"inference_candidates_ready",
|
||||
"phase_rerank_started",
|
||||
"phase_rerank_candidate_started",
|
||||
"phase_rerank_candidate_step",
|
||||
"phase_rerank_candidate_completed",
|
||||
"phase_selection_completed",
|
||||
"analysis_completed",
|
||||
"phase_field_prepared",
|
||||
|
|
@ -86,6 +91,7 @@ const state = {
|
|||
guidanceImage: null,
|
||||
inference: null,
|
||||
inferenceMode: null,
|
||||
latticeSearch: null,
|
||||
phaseRerank: null,
|
||||
phaseFieldPrep: null,
|
||||
frames: [],
|
||||
|
|
@ -178,6 +184,7 @@ function resetRunArtifacts(options = {}) {
|
|||
state.guidanceImage = null;
|
||||
state.inference = null;
|
||||
state.inferenceMode = null;
|
||||
state.latticeSearch = null;
|
||||
state.phaseRerank = null;
|
||||
state.phaseFieldPrep = null;
|
||||
state.frames = [];
|
||||
|
|
@ -322,7 +329,22 @@ function renderStatusMetrics() {
|
|||
statusMetrics.innerHTML = "";
|
||||
const items = [];
|
||||
const frame = getSelectedFrame();
|
||||
if (state.stageKey === "inference" || state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.stageKey === "inference") {
|
||||
if (state.latticeSearch) {
|
||||
items.push(["Sizes", `${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}`]);
|
||||
items.push(["Phase samples", String(state.latticeSearch.phaseSampleCount)]);
|
||||
if (state.latticeSearch.currentTargetWidth && state.latticeSearch.currentTargetHeight) {
|
||||
items.push(["Current grid", `${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.latticeSearch.device) {
|
||||
items.push(["Device", state.latticeSearch.device]);
|
||||
}
|
||||
if (state.latticeSearch.bestScore !== null) {
|
||||
items.push(["Best score", formatNumber(state.latticeSearch.bestScore, 3)]);
|
||||
}
|
||||
}
|
||||
}
|
||||
else if (state.stageKey === "analysis" || state.stageKey === "selection") {
|
||||
if (state.inference) {
|
||||
items.push(["Mode", state.inferenceMode === "fixed" ? "pinned" : state.inferenceMode ?? "search"]);
|
||||
items.push(["Grid", `${state.inference.target_width} x ${state.inference.target_height}`]);
|
||||
|
|
@ -332,8 +354,25 @@ function renderStatusMetrics() {
|
|||
}
|
||||
else if (state.stageKey === "rerank") {
|
||||
if (state.phaseRerank) {
|
||||
items.push(["Preview steps", String(state.phaseRerank.previewSteps)]);
|
||||
items.push(["Candidates", String(state.phaseRerank.candidateCount)]);
|
||||
items.push(["Candidates", `${state.phaseRerank.completedCandidates} / ${state.phaseRerank.candidateCount}`]);
|
||||
if (state.phaseRerank.currentCandidateIndex > 0) {
|
||||
items.push(["Active candidate", `${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}`]);
|
||||
}
|
||||
items.push([
|
||||
"Preview step",
|
||||
state.phaseRerank.currentTotalSteps > 0
|
||||
? `${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps}`
|
||||
: String(state.phaseRerank.previewSteps),
|
||||
]);
|
||||
if (state.phaseRerank.currentTargetWidth && state.phaseRerank.currentTargetHeight) {
|
||||
items.push(["Grid", `${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentPhaseX !== null && state.phaseRerank.currentPhaseY !== null) {
|
||||
items.push(["Phase", `${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}`]);
|
||||
}
|
||||
if (state.phaseRerank.currentLoss !== null) {
|
||||
items.push(["Loss", formatNumber(state.phaseRerank.currentLoss, 4)]);
|
||||
}
|
||||
items.push(["Confidence", formatNumber(state.phaseRerank.confidence, 3)]);
|
||||
}
|
||||
}
|
||||
|
|
@ -745,6 +784,40 @@ async function handleEvent(eventName, payload) {
|
|||
state.preprocessedImage = payload.sourceImage;
|
||||
addLog("Preprocess", "Stripped edge-connected background noise.");
|
||||
break;
|
||||
case "lattice_search_started":
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: typeof payload.device === "string" ? payload.device : null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
setStage("inference", "Lattice search", `Testing ${state.latticeSearch.candidateCount} size candidates with ${state.latticeSearch.phaseSampleCount} phase samples each.`);
|
||||
addLog("Search", `Started lattice search across ${String(payload.candidateCount)} size candidates.`);
|
||||
break;
|
||||
case "lattice_search_progress":
|
||||
if (!state.latticeSearch) {
|
||||
state.latticeSearch = {
|
||||
candidateCount: Number(payload.totalCandidates ?? 0),
|
||||
completedCandidates: 0,
|
||||
phaseSampleCount: Number(payload.phaseSampleCount ?? 0),
|
||||
device: null,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
bestScore: null,
|
||||
};
|
||||
}
|
||||
state.latticeSearch.completedCandidates = Number(payload.completedCandidates ?? state.latticeSearch.completedCandidates);
|
||||
state.latticeSearch.candidateCount = Number(payload.totalCandidates ?? state.latticeSearch.candidateCount);
|
||||
state.latticeSearch.phaseSampleCount = Number(payload.phaseSampleCount ?? state.latticeSearch.phaseSampleCount);
|
||||
state.latticeSearch.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.latticeSearch.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.latticeSearch.bestScore =
|
||||
payload.bestScore === null || payload.bestScore === undefined ? state.latticeSearch.bestScore : Number(payload.bestScore);
|
||||
setStage("inference", "Lattice search", `Size ${state.latticeSearch.completedCandidates} / ${state.latticeSearch.candidateCount}: scored ${state.latticeSearch.currentTargetWidth} x ${state.latticeSearch.currentTargetHeight} over ${state.latticeSearch.phaseSampleCount} phase samples.`);
|
||||
break;
|
||||
case "inference_candidates_ready":
|
||||
state.inference = payload.inference;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
@ -755,10 +828,63 @@ async function handleEvent(eventName, payload) {
|
|||
previewSteps: Number(payload.previewSteps ?? 0),
|
||||
candidateCount: Number(payload.candidateCount ?? 0),
|
||||
confidence: Number(payload.confidence ?? 0),
|
||||
completedCandidates: 0,
|
||||
currentCandidateIndex: 0,
|
||||
currentTargetWidth: null,
|
||||
currentTargetHeight: null,
|
||||
currentPhaseX: null,
|
||||
currentPhaseY: null,
|
||||
currentStep: 0,
|
||||
currentTotalSteps: 0,
|
||||
currentLoss: null,
|
||||
};
|
||||
setStage("rerank", "Phase rerank", `Previewing ${state.phaseRerank.candidateCount} low-confidence candidates over ${state.phaseRerank.previewSteps} short solver steps.`);
|
||||
addLog("Rerank", `Running ${String(payload.previewSteps)} preview steps across ${String(payload.candidateCount)} low-confidence candidates.`);
|
||||
break;
|
||||
case "phase_rerank_candidate_started":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = 0;
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.previewSteps ?? state.phaseRerank.previewSteps);
|
||||
state.phaseRerank.currentLoss = null;
|
||||
setStage("rerank", "Phase rerank", `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: previewing ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight} at phase ${formatNumber(state.phaseRerank.currentPhaseX, 2)}, ${formatNumber(state.phaseRerank.currentPhaseY, 2)}.`);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_step":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.currentStep = Number(payload.step ?? 0);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? 0);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.loss === null || payload.loss === undefined ? null : Number(payload.loss);
|
||||
setStage("rerank", "Phase rerank", state.phaseRerank.currentTotalSteps <= 0
|
||||
? `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}: preview solve skipped.`
|
||||
: `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount}, preview step ${state.phaseRerank.currentStep} / ${state.phaseRerank.currentTotalSteps} on ${state.phaseRerank.currentTargetWidth} x ${state.phaseRerank.currentTargetHeight}.`);
|
||||
}
|
||||
break;
|
||||
case "phase_rerank_candidate_completed":
|
||||
if (state.phaseRerank) {
|
||||
state.phaseRerank.currentCandidateIndex = Number(payload.candidateIndex ?? state.phaseRerank.currentCandidateIndex);
|
||||
state.phaseRerank.currentTargetWidth = Number(payload.targetWidth ?? 0);
|
||||
state.phaseRerank.currentTargetHeight = Number(payload.targetHeight ?? 0);
|
||||
state.phaseRerank.currentPhaseX = Number(payload.phaseX ?? 0);
|
||||
state.phaseRerank.currentPhaseY = Number(payload.phaseY ?? 0);
|
||||
state.phaseRerank.completedCandidates = Number(payload.completedCandidates ?? state.phaseRerank.completedCandidates);
|
||||
state.phaseRerank.currentStep = Number(payload.totalSteps ?? state.phaseRerank.currentStep);
|
||||
state.phaseRerank.currentTotalSteps = Number(payload.totalSteps ?? state.phaseRerank.currentTotalSteps);
|
||||
state.phaseRerank.currentLoss =
|
||||
payload.finalLoss === null || payload.finalLoss === undefined ? state.phaseRerank.currentLoss : Number(payload.finalLoss);
|
||||
setStage("rerank", "Phase rerank", `Candidate ${state.phaseRerank.currentCandidateIndex} / ${state.phaseRerank.candidateCount} scored. ${state.phaseRerank.completedCandidates} of ${state.phaseRerank.candidateCount} previews finished.`);
|
||||
}
|
||||
break;
|
||||
case "phase_selection_completed":
|
||||
state.inference = payload.inference;
|
||||
state.inferenceMode = typeof payload.inferenceMode === "string" ? payload.inferenceMode : state.inferenceMode;
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import numpy as np
|
|||
|
||||
from .io import premultiply
|
||||
from .metrics import source_lattice_evidence_breakdown
|
||||
from .observe import PipelineObserver, emit_observer
|
||||
from .types import InferenceCandidate, InferenceResult
|
||||
|
||||
|
||||
|
|
@ -471,7 +472,12 @@ def _score_phase_group(
|
|||
return candidates
|
||||
|
||||
|
||||
def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str = "auto") -> InferenceResult:
|
||||
def infer_lattice(
|
||||
rgba: np.ndarray,
|
||||
target_size: int | None = None,
|
||||
device: str = "auto",
|
||||
observer: PipelineObserver | None = None,
|
||||
) -> InferenceResult:
|
||||
torch, _ = _require_torch()
|
||||
resolved_device = _resolve_device(torch, device)
|
||||
height, width = rgba.shape[:2]
|
||||
|
|
@ -479,21 +485,40 @@ def infer_lattice(rgba: np.ndarray, target_size: int | None = None, device: str
|
|||
hinted_sizes = _hint_target_sizes_from_spacing(width, height, spacing_x, spacing_y)
|
||||
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)
|
||||
candidate_dims = _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes)
|
||||
phase_sample_count = int(phase_values.size * phase_values.size)
|
||||
|
||||
emit_observer(
|
||||
observer,
|
||||
"lattice_search_started",
|
||||
candidate_count=len(candidate_dims),
|
||||
phase_sample_count=phase_sample_count,
|
||||
device=resolved_device,
|
||||
)
|
||||
|
||||
candidates: list[InferenceCandidate] = []
|
||||
for target_width, target_height in _candidate_dims(width, height, target_size, hinted_sizes=hinted_sizes):
|
||||
candidates.extend(
|
||||
_score_phase_group(
|
||||
rgba,
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
prior_cell_x=prior_cell_x,
|
||||
prior_cell_y=prior_cell_y,
|
||||
prior_reliability=prior_reliability,
|
||||
phase_x_values=phase_values,
|
||||
phase_y_values=phase_values,
|
||||
device=resolved_device,
|
||||
)
|
||||
for candidate_index, (target_width, target_height) in enumerate(candidate_dims, start=1):
|
||||
scored_group = _score_phase_group(
|
||||
rgba,
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
prior_cell_x=prior_cell_x,
|
||||
prior_cell_y=prior_cell_y,
|
||||
prior_reliability=prior_reliability,
|
||||
phase_x_values=phase_values,
|
||||
phase_y_values=phase_values,
|
||||
device=resolved_device,
|
||||
)
|
||||
candidates.extend(scored_group)
|
||||
emit_observer(
|
||||
observer,
|
||||
"lattice_search_progress",
|
||||
completed_candidates=candidate_index,
|
||||
total_candidates=len(candidate_dims),
|
||||
target_width=int(target_width),
|
||||
target_height=int(target_height),
|
||||
phase_sample_count=phase_sample_count,
|
||||
best_score=None if not scored_group else float(max(candidate.score for candidate in scored_group)),
|
||||
)
|
||||
|
||||
candidates.sort(key=lambda item: item.score, reverse=True)
|
||||
|
|
|
|||
|
|
@ -125,7 +125,7 @@ def run_pipeline_rgba(
|
|||
label="Lattice search",
|
||||
detail="Searching for the output size and phase that best fit the input.",
|
||||
)
|
||||
inference = infer_lattice(source, target_size=target_size, device=device)
|
||||
inference = infer_lattice(source, target_size=target_size, device=device, observer=observer)
|
||||
inference_mode = "searched"
|
||||
else:
|
||||
emit_observer(
|
||||
|
|
@ -373,8 +373,9 @@ def _select_phase_candidate_with_reconstruction(
|
|||
candidate_count=min(8, len(inference.top_candidates)),
|
||||
confidence=float(inference.confidence),
|
||||
)
|
||||
total_candidates = min(8, len(inference.top_candidates))
|
||||
candidate_records: list[dict[str, float | InferenceResult]] = []
|
||||
for candidate in inference.top_candidates[:8]:
|
||||
for candidate_index, candidate in enumerate(inference.top_candidates[:total_candidates], start=1):
|
||||
candidate_inference = InferenceResult(
|
||||
target_width=candidate.target_width,
|
||||
target_height=candidate.target_height,
|
||||
|
|
@ -383,6 +384,39 @@ def _select_phase_candidate_with_reconstruction(
|
|||
confidence=inference.confidence,
|
||||
top_candidates=inference.top_candidates,
|
||||
)
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_started",
|
||||
candidate_index=candidate_index,
|
||||
total_candidates=total_candidates,
|
||||
target_width=int(candidate.target_width),
|
||||
target_height=int(candidate.target_height),
|
||||
phase_x=float(candidate.phase_x),
|
||||
phase_y=float(candidate.phase_y),
|
||||
preview_steps=preview_steps,
|
||||
)
|
||||
|
||||
def rerank_observer(event: str, payload: dict[str, Any]) -> None:
|
||||
if observer is None:
|
||||
return
|
||||
common = {
|
||||
"candidate_index": candidate_index,
|
||||
"total_candidates": total_candidates,
|
||||
"target_width": int(candidate.target_width),
|
||||
"target_height": int(candidate.target_height),
|
||||
"phase_x": float(candidate.phase_x),
|
||||
"phase_y": float(candidate.phase_y),
|
||||
}
|
||||
if event in {"phase_field_initial", "phase_field_step"}:
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_step",
|
||||
step=int(payload["step"]),
|
||||
total_steps=int(payload["total_steps"]),
|
||||
loss=None if payload.get("loss") is None else float(payload["loss"]),
|
||||
**common,
|
||||
)
|
||||
|
||||
candidate_artifacts, _candidate_diagnostics = _run_reconstruction(
|
||||
source,
|
||||
inference=candidate_inference,
|
||||
|
|
@ -391,7 +425,22 @@ def _select_phase_candidate_with_reconstruction(
|
|||
seed=seed,
|
||||
device=device,
|
||||
solver_params=solver_params,
|
||||
observer=None,
|
||||
observer=rerank_observer if observer is not None else None,
|
||||
)
|
||||
emit_observer(
|
||||
observer,
|
||||
"phase_rerank_candidate_completed",
|
||||
candidate_index=candidate_index,
|
||||
total_candidates=total_candidates,
|
||||
completed_candidates=candidate_index,
|
||||
target_width=int(candidate.target_width),
|
||||
target_height=int(candidate.target_height),
|
||||
phase_x=float(candidate.phase_x),
|
||||
phase_y=float(candidate.phase_y),
|
||||
total_steps=preview_steps,
|
||||
final_loss=(
|
||||
None if not candidate_artifacts.loss_history else float(candidate_artifacts.loss_history[-1])
|
||||
),
|
||||
)
|
||||
support = source_lattice_consistency_breakdown(
|
||||
source,
|
||||
|
|
|
|||
Loading…
Reference in a new issue