diff --git a/extra/viz/cli.py b/extra/viz/cli.py
index faf481452d..74bb5b4269 100755
--- a/extra/viz/cli.py
+++ b/extra/viz/cli.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
-import argparse, pathlib, signal, sys, struct, json, itertools
+import argparse, pathlib, signal, sys, struct, json, itertools, os
+os.environ["VIZ"] = "0"
 if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL)
 from typing import Iterator
 from tinygrad.viz import serve as viz
@@ -24,9 +25,8 @@ def decode_profile(data:bytes) -> dict:
     klen = u("<B")[0]
     k = ret[off:off+klen].decode()
     off += klen
-    v:dict = {"events":[]}
-    layout[k] = v
     event_type, event_count = u("<BI")
+    layout[k] = v = {"event_type":event_type, "events":[]}
     if event_type == 0:
       for _ in range(event_count):
         name, ref, key, st, dur, fmt = u("<IIIIfI")
@@ -41,7 +41,8 @@ def decode_profile(data:bytes) -> dict:
         else:
           alloc, ts, key = u("<BII")
           if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[(k, strings[rep], num, mode) \
+              for k,rep,num,mode in [u("<IIIB") for _ in range(u("<I")[0])]]}})
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 def get(data:dict, key:str):
@@ -114,6 +115,17 @@ def main(args) -> None:
       print(tabulate(rows, headers=cols, tablefmt="github"))
       return None
 
+    # ** Memory printer
+    if data["event_type"] == 1 and data.get("events", []):
+      print(f"Peak: {data['peak']}"+"\n"+f"{'TS':<10}  {'Event':<6}  {'Key':>8}  Info")
+      modes = ("read","write","write+read")
+      for e in data["events"]:
+        info = str(e.get("arg", {}))
+        if e["event"] == "free":
+          info = ', '.join([f"{format_colored(kernel)} {['read','write','write+read'][mode]}@data{num}" for _,kernel,num,mode in e["arg"]["users"]])
+        print(f"{e['ts']:<10}  {e['event']:<6}  {e.get('key', ''):>8}  {info}")
+      return None
+
     # ** Profiler printer
     agg:dict[str, tuple[float, int]] = {}
     total = 0
diff --git a/test/null/test_viz.py b/test/null/test_viz.py
index 3b94ee034e..e98f6b4458 100644
--- a/test/null/test_viz.py
+++ b/test/null/test_viz.py
@@ -2,11 +2,11 @@ import unittest, decimal, sys, json, contextlib
 from dataclasses import dataclass
 from typing import Generator
 
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS, profile_matches
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
-from tinygrad.helpers import VIZ, cpu_profile, ProfilePointEvent, unwrap
+from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
+from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
 from tinygrad.device import Buffer
 
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
@@ -37,25 +37,13 @@ class VizTrace:
 
 @contextlib.contextmanager
 def save_viz():
-  # clear previous traces
   for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
   Buffer.profile_events.clear()
   cpu_events.clear()
-  # set the context vars to enable VIZ
-  prev_viz = VIZ.value
-  VIZ.value = -1
-  prev_tms = TRACK_MATCH_STATS.value
-  TRACK_MATCH_STATS.value = 2
-  prev_profile = PROFILE.value
-  PROFILE.value = 1
   viz = VizTrace()
-  try:
+  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1):
     yield viz
-  finally:
-    viz.set_data()
-    TRACK_MATCH_STATS.value = prev_tms
-    PROFILE.value = prev_profile
-    VIZ.value = prev_viz
+  viz.set_data()
 
 class TestViz(unittest.TestCase):
   def test_simple(self):