From 5a56f1fa1eac803fd135cad0e4afcdba4374cdaf Mon Sep 17 00:00:00 2001
From: Meta Vens <rwvens@hotmail.com>
Date: Sat, 25 Apr 2026 03:46:06 +0100
Subject: [PATCH] Reset viewer after debug stages

---
 frontend/src/app.ts               | 45 ++++++++++++++++-------------
 src/repixelizer/gui_static/app.js | 47 +++++++++++++++++--------------
 2 files changed, 51 insertions(+), 41 deletions(-)

diff --git a/frontend/src/app.ts b/frontend/src/app.ts
index 5914405..c04b2a3 100644
--- a/frontend/src/app.ts
+++ b/frontend/src/app.ts
@@ -595,31 +595,36 @@ function renderSummary(): void {
 
 async function renderViewer(): Promise<void> {
   const frame = getSelectedFrame();
-  let leftAsset = state.preprocessedImage ?? state.sourceImage;
+  const debugViewerStage = new Set(["inference", "analysis", "selection", "rerank", "solver", "cleanup"]);
+  const showDebugViewer = debugViewerStage.has(state.stageKey);
+  let leftAsset = state.sourceImage;
   let rightAsset = state.finalOutputImage ?? state.cleanupImage ?? state.sourceImage;
   let leftLabel = "Input";
   let rightLabel = "Output";
 
-  if (state.latticeImage) {
-    leftAsset = state.latticeImage;
-    leftLabel = "Lattice Prep";
-  }
-  if (state.guidanceImage) {
-    leftAsset = state.guidanceImage;
-    leftLabel = "Guidance";
-  }
-  if (frame) {
-    leftAsset = frame.samplingOverlayImage;
-    rightAsset = frame.outputImage;
-    leftLabel = "Sampling Overlay";
-    rightLabel = "Current Output";
-  } else if (state.cleanupImage) {
-    if (state.heatmapImage) {
-      leftAsset = state.heatmapImage;
-      leftLabel = "Cleanup Heatmap";
+  if (showDebugViewer) {
+    leftAsset = state.preprocessedImage ?? state.sourceImage;
+    if (state.latticeImage) {
+      leftAsset = state.latticeImage;
+      leftLabel = "Lattice Prep";
+    }
+    if (state.guidanceImage) {
+      leftAsset = state.guidanceImage;
+      leftLabel = "Guidance";
+    }
+    if (frame) {
+      leftAsset = frame.samplingOverlayImage;
+      rightAsset = frame.outputImage;
+      leftLabel = "Sampling Overlay";
+      rightLabel = "Current Output";
+    } else if (state.cleanupImage) {
+      if (state.heatmapImage) {
+        leftAsset = state.heatmapImage;
+        leftLabel = "Cleanup Heatmap";
+      }
+      rightAsset = state.cleanupImage;
+      rightLabel = "Cleaned Output";
     }
-    rightAsset = state.cleanupImage;
-    rightLabel = "Cleaned Output";
   }
 
   leftVizLabel.textContent = leftLabel;
diff --git a/src/repixelizer/gui_static/app.js b/src/repixelizer/gui_static/app.js
index e1facc7..ea74f69 100644
--- a/src/repixelizer/gui_static/app.js
+++ b/src/repixelizer/gui_static/app.js
@@ -449,31 +449,36 @@ function renderSummary() {
 }
 async function renderViewer() {
     const frame = getSelectedFrame();
-    let leftAsset = state.preprocessedImage ?? state.sourceImage;
+    const debugViewerStage = new Set(["inference", "analysis", "selection", "rerank", "solver", "cleanup"]);
+    const showDebugViewer = debugViewerStage.has(state.stageKey);
+    let leftAsset = state.sourceImage;
     let rightAsset = state.finalOutputImage ?? state.cleanupImage ?? state.sourceImage;
     let leftLabel = "Input";
     let rightLabel = "Output";
-    if (state.latticeImage) {
-        leftAsset = state.latticeImage;
-        leftLabel = "Lattice Prep";
-    }
-    if (state.guidanceImage) {
-        leftAsset = state.guidanceImage;
-        leftLabel = "Guidance";
-    }
-    if (frame) {
-        leftAsset = frame.samplingOverlayImage;
-        rightAsset = frame.outputImage;
-        leftLabel = "Sampling Overlay";
-        rightLabel = "Current Output";
-    }
-    else if (state.cleanupImage) {
-        if (state.heatmapImage) {
-            leftAsset = state.heatmapImage;
-            leftLabel = "Cleanup Heatmap";
+    if (showDebugViewer) {
+        leftAsset = state.preprocessedImage ?? state.sourceImage;
+        if (state.latticeImage) {
+            leftAsset = state.latticeImage;
+            leftLabel = "Lattice Prep";
+        }
+        if (state.guidanceImage) {
+            leftAsset = state.guidanceImage;
+            leftLabel = "Guidance";
+        }
+        if (frame) {
+            leftAsset = frame.samplingOverlayImage;
+            rightAsset = frame.outputImage;
+            leftLabel = "Sampling Overlay";
+            rightLabel = "Current Output";
+        }
+        else if (state.cleanupImage) {
+            if (state.heatmapImage) {
+                leftAsset = state.heatmapImage;
+                leftLabel = "Cleanup Heatmap";
+            }
+            rightAsset = state.cleanupImage;
+            rightLabel = "Cleaned Output";
         }
-        rightAsset = state.cleanupImage;
-        rightLabel = "Cleaned Output";
     }
     leftVizLabel.textContent = leftLabel;
     rightVizLabel.textContent = rightLabel;