From b4ea45b4a656290b5e3848cfc923255c5825efc3 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 28 Mar 2025 15:24:53 +0800
Subject: [PATCH] fix viz recenter + worker cleanup (#9607)

---
 tinygrad/viz/index.html   |  5 +++--
 tinygrad/viz/lib/graph.js | 26 +++++++++++---------------
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html
index 8f33163715..606f820c60 100644
--- a/tinygrad/viz/index.html
+++ b/tinygrad/viz/index.html
@@ -262,7 +262,7 @@
     const x0 = document.querySelector(".kernel-list-parent").getBoundingClientRect().right;
     const x1 = document.querySelector(".metadata").getBoundingClientRect().left;
     const pad = 16;
-    const R = { x: x0+pad, y: mainRect.top+pad, width: x1-x0-2*pad, height: mainRect.height-2*pad };
+    const R = { x: x0+pad, y: mainRect.top+pad, width: (x1>0 ? x1-x0 : mainRect.width)-2*pad, height: mainRect.height-2*pad };
     const r = document.querySelector("#render").getBoundingClientRect();
     if (r.width === 0) return;
     const scale = Math.min(R.width/r.width, R.height/r.height);
@@ -416,7 +416,8 @@
   document.querySelector(".collapse-btn").addEventListener("click", (e) => {
     isCollapsed = !isCollapsed;
     mainContainer.classList.toggle("collapsed", isCollapsed);
-    e.target.style.transform = isCollapsed ? "rotate(180deg)" : "rotate(0deg)";
+    e.currentTarget.blur();
+    e.currentTarget.style.transform = isCollapsed ? "rotate(180deg)" : "rotate(0deg)";
   });
   // **** resizer
   function appendResizer(element, { minWidth, maxWidth }, left=false) {
diff --git a/tinygrad/viz/lib/graph.js b/tinygrad/viz/lib/graph.js
index c3c08a69dd..e8e09b11b5 100644
--- a/tinygrad/viz/lib/graph.js
+++ b/tinygrad/viz/lib/graph.js
@@ -8,31 +8,27 @@ function intersectRect(r1, r2) {
   return {x:r1.x+dx*scale, y:r1.y+dy*scale};
 }
 
-const allWorkers = [];
-let workerUrl = null;
+let [workerUrl, worker, timeout] = [null, null, null];
 window.renderGraph = async function(graph, additions, name) {
-  if (workerUrl == null) {
-    const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/lib/worker.js"].map(u => fetch(u)));
-    workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
-  }
-  while (allWorkers.length) {
-    const { worker, timeout } = allWorkers.pop();
-    worker.terminate();
-    clearTimeout(timeout);
-  }
-
   if (name === "View Memory Graph") {
     return renderMemoryGraph(graph);
   }
   d3.select("#bars").html("");
 
   // ** start calculating the new layout (non-blocking)
-  worker = new Worker(workerUrl);
+  if (worker == null) {
+    const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/lib/worker.js"].map(u => fetch(u)));
+    workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
+    worker = new Worker(workerUrl);
+  } else {
+    worker.terminate();
+    worker = new Worker(workerUrl);
+  }
+  if (timeout != null) clearTimeout(timeout);
   const progressMessage = document.querySelector(".progress-message");
-  const timeout = setTimeout(() => {
+  timeout = setTimeout(() => {
     progressMessage.style.display = "block";
   }, 2000);
-  allWorkers.push({worker, timeout});
   worker.postMessage({graph, additions});
 
   worker.onmessage = (e) => {