update to use gradients

app.py

@@ -10,6 +10,7 @@ import random
 import beartype
 import einops.layers.torch
 import gradio as gr
+import matplotlib
 import numpy as np
 import open_clip
 import requests
@@ -18,7 +19,7 @@ import torch
 from jaxtyping import Float, jaxtyped
 from PIL import Image, ImageDraw
 from torch import Tensor
-from torchvision.transforms import v2
+from torchvision import transforms
 
 log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=log_format)
@@ -29,10 +30,10 @@ logger = logging.getLogger("app.py")
 # Global Constants #
 ####################
 
-DEBUG = True
+DEBUG = False
 """Whether we are debugging."""
 
-n_sae_latents = 3
+n_sae_latents = 5
 """Number of SAE latents to show."""
 
 n_sae_examples = 4
@@ -54,6 +55,7 @@ CWD = pathlib.Path(__file__).parent
 
 r2_url = "https://pub-289086e849214430853bc87bd8964988.r2.dev/"
 
+colormap = matplotlib.colormaps.get_cmap("plasma")
 
 logger.info("Set global constants.")
 
@@ -117,7 +119,6 @@ def make_img(
         (resize_w_px + crop_w_px) // 2,
         (resize_h_px + crop_h_px) // 2,
     )
-
     img = img.resize(resize_size_px).crop(crop_coords_px)
     img = add_highlights(img, patches.numpy(), upper=upper, opacity=0.5)
     return img
@@ -207,10 +208,10 @@ logger.info("Loaded SAE.")
 # Datasets #
 ############
 
-human_transform = v2.Compose([
-    v2.Resize((512, 512), interpolation=v2.InterpolationMode.NEAREST),
-    v2.CenterCrop((448, 448)),
-    v2.ToImage(),
+human_transform = transforms.Compose([
+    transforms.Resize((448,), interpolation=transforms.InterpolationMode.BICUBIC),
+    transforms.CenterCrop((448, 448)),
+    transforms.ToTensor(),
     einops.layers.torch.Rearrange("channels width height -> width height channels"),
 ])
 
@@ -258,7 +259,10 @@ mask = mask & (sparsity < max_frequency)
 def get_image(image_i: int) -> list[Image.Image | int]:
     image = get_dataset_img(image_i)
     image = human_transform(image)
-    return [Image.fromarray(image.numpy()), image_labels[image_i]]
+    return [
+        Image.fromarray((image * 255).to(torch.uint8).numpy()),
+        image_labels[image_i],
+    ]
 
 
 @beartype.beartype
@@ -268,7 +272,7 @@ def get_random_class_image(cls: int) -> Image.Image:
 
     image = get_dataset_img(i)
     image = human_transform(image)
-    return Image.fromarray(image.numpy())
+    return Image.fromarray((image * 255).to(torch.uint8).numpy())
 
 
 @torch.inference_mode
@@ -430,8 +434,9 @@ def add_highlights(
     overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
     draw = ImageDraw.Draw(overlay)
 
-    # Using semi-transparent red (255, 0, 0, alpha)
-    for p, val in enumerate(patches):
+    colors = (colormap(patches / (upper + 1e-9))[:, :3] * 255).astype(np.uint8)
+
+    for p, (val, color) in enumerate(zip(patches, colors)):
         assert upper is not None
         val /= upper + 1e-9
         x_np, y_np = p % iw_np, p // ih_np
@@ -440,7 +445,7 @@ def add_highlights(
                 (x_np * pw_px, y_np * ph_px),
                 (x_np * pw_px + pw_px, y_np * ph_px + ph_px),
             ],
-            fill=(int(val * 256), 0, 0, int(opacity * val * 256)),
+            fill=(*color, int(opacity * val * 255)),
         )
 
     # Composite the original image and the overlay

requirements.txt

@@ -2,9 +2,9 @@
 #    uv pip compile pyproject.toml
 aiofiles==23.2.1
     # via gradio
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.12
     # via
     #   datasets
     #   fsspec
@@ -18,43 +18,22 @@ anyio==4.8.0
     # via
     #   gradio
     #   httpx
-    #   jupyter-server
+    #   pycrdt
     #   starlette
-argon2-cffi==23.1.0
-    # via jupyter-server
-argon2-cffi-bindings==21.2.0
-    # via argon2-cffi
-arrow==1.3.0
-    # via isoduration
-asttokens==3.0.0
-    # via stack-data
-async-lru==2.0.4
-    # via jupyterlab
-attrs==24.3.0
+attrs==25.1.0
     # via
     #   aiohttp
     #   jsonschema
     #   referencing
-babel==2.16.0
-    # via jupyterlab-server
 beartype==0.19.0
     # via
     #   saev-image-classification (pyproject.toml)
     #   saev
-beautifulsoup4==4.12.3
-    # via nbconvert
-bleach==6.2.0
-    # via nbconvert
-braceexpand==0.1.7
-    # via webdataset
-certifi==2024.12.14
+certifi==2025.1.31
     # via
     #   httpcore
     #   httpx
     #   requests
-    #   sentry-sdk
-cffi==1.17.1
-    # via argon2-cffi-bindings
 charset-normalizer==3.4.1
     # via requests
 click==8.1.8
@@ -62,60 +41,42 @@ click==8.1.8
     #   marimo
     #   typer
     #   uvicorn
-    #   wandb
-cloudpickle==3.1.0
+cloudpickle==3.1.1
     # via submitit
-comm==0.2.2
-    # via ipykernel
 contourpy==1.3.1
     # via matplotlib
 cycler==0.12.1
     # via matplotlib
-datasets==3.2.0
+datasets==3.3.0
     # via saev
-debugpy==1.8.11
-    # via ipykernel
-decorator==5.1.1
-    # via ipython
-defusedxml==0.7.1
-    # via nbconvert
 dill==0.3.8
     # via
     #   datasets
     #   multiprocess
-docker-pycreds==0.4.0
-    # via wandb
 docstring-parser==0.16
     # via tyro
 docutils==0.21.2
     # via marimo
-einops==0.8.0
+einops==0.8.1
     # via
     #   saev-image-classification (pyproject.toml)
     #   saev
-executing==2.1.0
-    # via stack-data
-fastapi==0.115.6
+fastapi==0.115.8
     # via gradio
-fastjsonschema==2.21.1
-    # via nbformat
 ffmpy==0.5.0
     # via gradio
-filelock==3.16.1
+filelock==3.17.0
     # via
     #   datasets
     #   huggingface-hub
     #   torch
-    #   triton
-fonttools==4.55.3
+fonttools==4.56.0
     # via matplotlib
-fqdn==1.5.1
-    # via jsonschema
 frozenlist==1.5.0
     # via
     #   aiohttp
     #   aiosignal
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
     #   datasets
     #   gradio-client
@@ -123,13 +84,9 @@ fsspec==2024.9.0
     #   torch
 ftfy==6.3.1
     # via open-clip-torch
-gitdb==4.0.12
-    # via gitpython
-gitpython==3.1.44
-    # via wandb
-gradio==5.10.0
+gradio==5.16.0
     # via saev-image-classification (pyproject.toml)
-gradio-client==1.5.3
+gradio-client==1.7.0
     # via gradio
 h11==0.14.0
     # via
@@ -141,9 +98,8 @@ httpx==0.28.1
     # via
     #   gradio
     #   gradio-client
-    #   jupyterlab
     #   safehttpx
-huggingface-hub==0.27.1
+huggingface-hub==0.28.1
     # via
     #   datasets
     #   gradio
@@ -154,90 +110,32 @@ idna==3.10
     # via
     #   anyio
     #   httpx
-    #   jsonschema
     #   requests
     #   yarl
-ipykernel==6.29.5
-    # via jupyterlab
-ipython==8.31.0
-    # via ipykernel
-isoduration==20.11.0
-    # via jsonschema
 itsdangerous==2.2.0
     # via marimo
-jaxtyping==0.2.36
+jaxtyping==0.2.38
     # via
     #   saev-image-classification (pyproject.toml)
     #   saev
 jedi==0.19.2
-    # via
-    #   ipython
-    #   marimo
+    # via marimo
 jinja2==3.1.5
     # via
     #   altair
     #   gradio
-    #   jupyter-server
-    #   jupyterlab
-    #   jupyterlab-server
-    #   nbconvert
     #   torch
-joblib==1.4.2
-    # via scikit-learn
-json5==0.10.0
-    # via jupyterlab-server
-jsonpointer==3.0.0
-    # via jsonschema
 jsonschema==4.23.0
-    # via
-    #   altair
-    #   jupyter-events
-    #   jupyterlab-server
-    #   nbformat
+    # via altair
 jsonschema-specifications==2024.10.1
     # via jsonschema
-jupyter-client==8.6.3
-    # via
-    #   ipykernel
-    #   jupyter-server
-    #   nbclient
-jupyter-core==5.7.2
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
-    #   jupyterlab
-    #   nbclient
-    #   nbconvert
-    #   nbformat
-jupyter-events==0.11.0
-    # via jupyter-server
-jupyter-lsp==2.2.5
-    # via jupyterlab
-jupyter-server==2.15.0
-    # via
-    #   jupyter-lsp
-    #   jupyterlab
-    #   jupyterlab-server
-    #   notebook-shim
-jupyter-server-terminals==0.5.3
-    # via jupyter-server
-jupyterlab==4.3.4
-    # via saev
-jupyterlab-pygments==0.3.0
-    # via nbconvert
-jupyterlab-server==2.27.3
-    # via jupyterlab
 kiwisolver==1.4.8
     # via matplotlib
-mako==1.3.8
-    # via pdoc3
-marimo==0.10.9
+marimo==0.11.5
     # via saev
 markdown==3.7
     # via
     #   marimo
-    #   pdoc3
     #   pymdown-extensions
 markdown-it-py==3.0.0
     # via rich
@@ -245,18 +143,10 @@ markupsafe==2.1.5
     # via
     #   gradio
     #   jinja2
-    #   mako
-    #   nbconvert
 matplotlib==3.10.0
     # via saev
-matplotlib-inline==0.1.7
-    # via
-    #   ipykernel
-    #   ipython
 mdurl==0.1.2
     # via markdown-it-py
-mistune==3.1.0
-    # via nbconvert
 mpmath==1.3.0
     # via sympy
 multidict==6.1.0
@@ -265,26 +155,13 @@ multidict==6.1.0
     #   yarl
 multiprocess==0.70.16
     # via datasets
-narwhals==1.21.1
+narwhals==1.26.0
     # via
     #   altair
     #   marimo
-nbclient==0.10.2
-    # via nbconvert
-nbconvert==7.16.5
-    # via jupyter-server
-nbformat==5.10.4
-    # via
-    #   jupyter-server
-    #   nbclient
-    #   nbconvert
-nest-asyncio==1.6.0
-    # via ipykernel
 networkx==3.4.2
     # via torch
-notebook-shim==0.2.4
-    # via jupyterlab
-numpy==2.2.1
+numpy==2.2.3
     # via
     #   saev-image-classification (pyproject.toml)
     #   contourpy
@@ -292,10 +169,7 @@ numpy==2.2.1
     #   gradio
     #   matplotlib
     #   pandas
-    #   scikit-learn
-    #   scipy
     #   torchvision
-    #   webdataset
 nvidia-cublas-cu12==12.4.5.8
     # via
     #   nvidia-cudnn-cu12
@@ -319,6 +193,8 @@ nvidia-cusparse-cu12==12.3.1.170
     # via
     #   nvidia-cusolver-cu12
     #   torch
+nvidia-cusparselt-cu12==0.6.2
+    # via torch
 nvidia-nccl-cu12==2.21.5
     # via torch
 nvidia-nvjitlink-cu12==12.4.127
@@ -332,10 +208,8 @@ open-clip-torch==2.30.0
     # via
     #   saev-image-classification (pyproject.toml)
     #   saev
-orjson==3.10.13
+orjson==3.10.15
     # via gradio
-overrides==7.7.0
-    # via jupyter-server
 packaging==24.2
     # via
     #   altair
@@ -343,130 +217,74 @@ packaging==24.2
     #   gradio
     #   gradio-client
     #   huggingface-hub
-    #   ipykernel
-    #   jupyter-server
-    #   jupyterlab
-    #   jupyterlab-server
     #   marimo
     #   matplotlib
-    #   nbconvert
 pandas==2.2.3
     # via
     #   datasets
     #   gradio
-pandocfilters==1.5.1
-    # via nbconvert
 parso==0.8.4
     # via jedi
-pdoc3==0.11.5
-    # via saev
-pexpect==4.9.0
-    # via ipython
 pillow==11.1.0
     # via
     #   gradio
     #   matplotlib
     #   saev
     #   torchvision
-platformdirs==4.3.6
-    # via
-    #   jupyter-core
-    #   wandb
-polars==1.19.0
+polars==1.22.0
     # via saev
-prometheus-client==0.21.1
-    # via jupyter-server
-prompt-toolkit==3.0.48
-    # via ipython
 propcache==0.2.1
     # via
     #   aiohttp
     #   yarl
-protobuf==5.29.2
-    # via wandb
-psutil==6.1.1
-    # via
-    #   ipykernel
-    #   marimo
-    #   wandb
-ptyprocess==0.7.0
-    # via
-    #   pexpect
-    #   terminado
-pure-eval==0.2.3
-    # via stack-data
-pyarrow==18.1.0
+psutil==7.0.0
+    # via marimo
+pyarrow==19.0.0
     # via datasets
-pycparser==2.22
-    # via cffi
-pydantic==2.10.4
+pycrdt==0.11.1
+    # via marimo
+pydantic==2.10.6
     # via
     #   fastapi
     #   gradio
-    #   wandb
 pydantic-core==2.27.2
     # via pydantic
 pydub==0.25.1
     # via gradio
 pygments==2.19.1
     # via
-    #   ipython
     #   marimo
-    #   nbconvert
     #   rich
-pymdown-extensions==10.14
+pymdown-extensions==10.14.3
     # via marimo
 pyparsing==3.2.1
     # via matplotlib
 python-dateutil==2.9.0.post0
     # via
-    #   arrow
-    #   jupyter-client
     #   matplotlib
     #   pandas
-python-json-logger==3.2.1
-    # via jupyter-events
 python-multipart==0.0.20
     # via gradio
-pytz==2024.2
+pytz==2025.1
     # via pandas
 pyyaml==6.0.2
     # via
     #   datasets
     #   gradio
     #   huggingface-hub
-    #   jupyter-events
     #   marimo
     #   pymdown-extensions
     #   timm
-    #   wandb
-    #   webdataset
-pyzmq==26.2.0
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
-referencing==0.35.1
+referencing==0.36.2
     # via
     #   jsonschema
     #   jsonschema-specifications
-    #   jupyter-events
 regex==2024.11.6
     # via open-clip-torch
 requests==2.32.3
     # via
     #   datasets
     #   huggingface-hub
-    #   jupyterlab-server
-    #   wandb
-rfc3339-validator==0.1.4
-    # via
-    #   jsonschema
-    #   jupyter-events
-rfc3986-validator==0.1.1
-    # via
-    #   jsonschema
-    #   jupyter-events
 rich==13.9.4
     # via
     #   typer
@@ -475,53 +293,31 @@ rpds-py==0.22.3
     # via
     #   jsonschema
     #   referencing
-ruff==0.8.6
+ruff==0.9.6
     # via
     #   gradio
     #   marimo
-saev @ git+https://github.com/samuelstevens/saev@c723ff95462736d907b3c1891d3a1496eb07318f
+saev @ git+https://github.com/samuelstevens/saev@83a442210c81b09e71e68e63ff7e5aadb84e0e87
     # via saev-image-classification (pyproject.toml)
 safehttpx==0.1.6
     # via gradio
-safetensors==0.5.1
+safetensors==0.5.2
     # via
     #   open-clip-torch
     #   timm
-scikit-learn==1.6.0
-    # via saev
-scipy==1.15.0
-    # via scikit-learn
 semantic-version==2.10.0
     # via gradio
-send2trash==1.8.3
-    # via jupyter-server
-sentry-sdk==2.19.2
-    # via wandb
-setproctitle==1.3.4
-    # via wandb
-setuptools==75.7.0
-    # via
-    #   jupyterlab
-    #   torch
-    #   wandb
+setuptools==75.8.0
+    # via torch
 shellingham==1.5.4
     # via typer
 shtab==1.7.1
     # via tyro
 six==1.17.0
-    # via
-    #   docker-pycreds
-    #   python-dateutil
-    #   rfc3339-validator
-smmap==5.0.2
-    # via gitdb
+    # via python-dateutil
 sniffio==1.3.1
     # via anyio
-soupsieve==2.6
-    # via beautifulsoup4
-stack-data==0.6.3
-    # via ipython
-starlette==0.41.3
+starlette==0.45.3
     # via
     #   fastapi
     #   gradio
@@ -530,67 +326,36 @@ submitit==1.5.2
     # via saev
 sympy==1.13.1
     # via torch
-terminado==0.18.1
-    # via
-    #   jupyter-server
-    #   jupyter-server-terminals
-threadpoolctl==3.5.0
-    # via scikit-learn
-timm==1.0.12
+timm==1.0.14
     # via open-clip-torch
-tinycss2==1.4.0
-    # via bleach
 tomlkit==0.13.2
     # via
     #   gradio
     #   marimo
-torch==2.5.1
+torch==2.6.0
     # via
     #   saev-image-classification (pyproject.toml)
     #   open-clip-torch
     #   saev
     #   timm
     #   torchvision
-torchvision==0.20.1
+torchvision==0.21.0
     # via
     #   saev-image-classification (pyproject.toml)
     #   open-clip-torch
     #   timm
-tornado==6.4.2
-    # via
-    #   ipykernel
-    #   jupyter-client
-    #   jupyter-server
-    #   jupyterlab
-    #   terminado
 tqdm==4.67.1
     # via
     #   datasets
     #   huggingface-hub
     #   open-clip-torch
     #   saev
-traitlets==5.14.3
-    # via
-    #   comm
-    #   ipykernel
-    #   ipython
-    #   jupyter-client
-    #   jupyter-core
-    #   jupyter-events
-    #   jupyter-server
-    #   jupyterlab
-    #   matplotlib-inline
-    #   nbclient
-    #   nbconvert
-    #   nbformat
-triton==3.1.0
+triton==3.2.0
     # via torch
 typeguard==4.4.1
     # via tyro
 typer==0.15.1
     # via gradio
-types-python-dateutil==2.9.0.20241206
-    # via arrow
 typing-extensions==4.12.2
     # via
     #   altair
@@ -601,44 +366,27 @@ typing-extensions==4.12.2
     #   huggingface-hub
     #   pydantic
     #   pydantic-core
+    #   referencing
     #   submitit
     #   torch
     #   typeguard
     #   typer
     #   tyro
-tyro==0.9.6
+tyro==0.9.14
     # via saev
-tzdata==2024.2
+tzdata==2025.1
     # via pandas
-uri-template==1.3.0
-    # via jsonschema
 urllib3==2.3.0
-    # via
-    #   requests
-    #   sentry-sdk
+    # via requests
 uvicorn==0.34.0
     # via
     #   gradio
     #   marimo
-vl-convert-python==1.7.0
-    # via saev
-wandb==0.19.1
-    # via saev
+wadler-lindig==0.1.3
+    # via jaxtyping
 wcwidth==0.2.13
-    # via
-    #   ftfy
-    #   prompt-toolkit
-webcolors==24.11.1
-    # via jsonschema
-webdataset==0.2.100
-    # via saev
-webencodings==0.5.1
-    # via
-    #   bleach
-    #   tinycss2
-websocket-client==1.8.0
-    # via jupyter-server
-websockets==14.1
+    # via ftfy
+websockets==14.2
     # via
     #   gradio-client
     #   marimo