initial commit: transformer x-ray (FastAPI + React, GIF export)

.dockerignore

@@ -0,0 +1,10 @@
+**/__pycache__
+**/.venv
+**/.env
+**/node_modules
+frontend/dist
+static
+.git
+.DS_Store
+*.log
+hf_cache

.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.pyc
+.venv/
+venv/
+.env
+node_modules/
+frontend/dist/
+static/
+.DS_Store
+.idea/
+.vscode/
+*.log
+.hf/
+hf_cache/

Dockerfile

@@ -0,0 +1,28 @@
+# ---- frontend build ----
+FROM node:20-alpine AS frontend
+WORKDIR /fe
+COPY frontend/package.json frontend/package-lock.json* ./
+RUN npm install --no-audit --no-fund
+COPY frontend/ ./
+RUN npm run build
+
+# ---- runtime ----
+FROM python:3.11-slim
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    HF_HOME=/tmp/hf \
+    TRANSFORMERS_CACHE=/tmp/hf \
+    MAX_PARAMS=1500000000
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git build-essential && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY backend/ ./backend/
+COPY --from=frontend /fe/dist ./static/
+
+EXPOSE 7860
+CMD ["uvicorn", "backend.main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,10 +1,43 @@
 ---
-title: Transformer Xray
-emoji: 🏢
-colorFrom: pink
-colorTo: red
+title: Transformer X-Ray
+emoji: 🔬
+colorFrom: indigo
+colorTo: purple
 sdk: docker
+app_port: 7860
 pinned: false
+license: mit
+short_description: Visualize any HF model layer-by-layer, export GIF
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Transformer X-Ray
+
+Interactive architecture visualizer for any model on the Hugging Face Hub.
+
+- Paste a model ID (e.g. `google/vit-base-patch16-224`, `bert-base-uncased`, `prajjwal1/bert-tiny`)
+- See every `nn.Module` from `transformers` as an interactive graph
+- Hit **Play** — real tensors from a real forward pass animate through the graph: tokens → embeddings → attention → MLP → logits → probabilities
+- Click any module to see its config, parameter count, and the actual tensor that flowed through it
+- Export the playthrough as a **GIF**
+
+## Architecture
+
+- **Backend** — FastAPI + `transformers` + `torch` (CPU). Loads the model from the Hub, traces it with `torch.fx` (with a `named_modules` fallback), registers forward hooks on every module, runs a sample input, and streams per-module tensor snapshots over a WebSocket.
+- **Frontend** — Vite + React + React Flow. Renders the graph, drives the play timeline, and assembles the GIF in-browser with `gif.js`.
+
+## Local dev
+
+```bash
+# backend
+pip install -r requirements.txt
+uvicorn backend.main:app --reload --port 7860
+
+# frontend (separate terminal)
+cd frontend && npm install && npm run dev
+```
+
+The Vite dev server proxies `/api` and `/ws` to the backend.
+
+## Deploy
+
+Push to a Hugging Face Space with `sdk: docker`. The included `Dockerfile` builds the frontend and serves it from the FastAPI app on port 7860.

backend/main.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import time
+import uuid
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from .model_loader import CACHE, load_model
+from .runner import get_last_run, run_model, store_run
+from .tracer import build_graph
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
+log = logging.getLogger("main")
+
+app = FastAPI(title="Transformer X-Ray")
+
+# Permissive CORS for local Vite dev. The deployed Space serves both apps from one origin.
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class LoadRequest(BaseModel):
+    model_id: str
+
+
+class RunRequest(BaseModel):
+    prompt: str | None = None
+
+
+def _config_summary(cfg: Any) -> dict:
+    """Pull common config fields. Falls back to to_dict() for everything else."""
+    full = cfg.to_dict() if hasattr(cfg, "to_dict") else {}
+    keys = (
+        "model_type", "architectures", "hidden_size", "num_hidden_layers",
+        "num_attention_heads", "num_key_value_heads", "intermediate_size",
+        "vocab_size", "max_position_embeddings", "image_size", "patch_size",
+        "num_channels", "id2label", "label2id", "torch_dtype", "tie_word_embeddings",
+    )
+    return {k: full.get(k) for k in keys if k in full}
+
+
+@app.post("/api/load")
+def api_load(req: LoadRequest):
+    try:
+        loaded = load_model(req.model_id)
+    except ValueError as e:
+        raise HTTPException(status_code=413, detail=str(e))
+    except Exception as e:
+        log.exception("load failed")
+        raise HTTPException(status_code=400, detail=f"failed to load: {e}")
+    tok = loaded.tokenizer
+    tokenizer_info = None
+    if tok is not None:
+        tokenizer_info = {
+            "name_or_path": getattr(tok, "name_or_path", None),
+            "vocab_size": getattr(tok, "vocab_size", None),
+            "model_max_length": getattr(tok, "model_max_length", None),
+            "special_tokens": list(getattr(tok, "all_special_tokens", []) or []),
+            "is_fast": getattr(tok, "is_fast", None),
+        }
+    return {
+        "model_id": loaded.model_id,
+        "modality": loaded.modality,
+        "head_kind": loaded.head_kind,
+        "param_count": loaded.param_count,
+        "config": _config_summary(loaded.config),
+        "tokenizer": tokenizer_info,
+        "processor": type(loaded.processor).__name__ if loaded.processor else None,
+    }
+
+
+@app.get("/api/graph")
+def api_graph():
+    loaded = CACHE.current()
+    if loaded is None:
+        raise HTTPException(status_code=400, detail="no model loaded")
+    return build_graph(loaded.model)
+
+
+@app.post("/api/run")
+def api_run(req: RunRequest):
+    loaded = CACHE.current()
+    if loaded is None:
+        raise HTTPException(status_code=400, detail="no model loaded")
+    t0 = time.time()
+    result = run_model(loaded, req.prompt)
+    elapsed = time.time() - t0
+    store_run(result)
+    run_id = uuid.uuid4().hex[:12]
+    return {"run_id": run_id, "steps": len(result.events), "elapsed_s": round(elapsed, 3)}
+
+
+@app.websocket("/ws/run/{run_id}")
+async def ws_run(ws: WebSocket, run_id: str):
+    await ws.accept()
+    result = get_last_run()
+    if result is None:
+        await ws.send_json({"error": "no run available"})
+        await ws.close()
+        return
+    try:
+        for ev in result.events:
+            await ws.send_json({
+                "step": ev.step,
+                "path": ev.path,
+                "module_class": ev.module_class,
+                "kind": ev.kind,
+                "payload": ev.payload,
+            })
+        await ws.send_json({"kind": "stream_end"})
+    except WebSocketDisconnect:
+        return
+
+
+@app.get("/api/health")
+def health():
+    return {"ok": True}
+
+
+# ---- static frontend ----
+STATIC_DIR = Path(os.environ.get("STATIC_DIR", "static"))
+if STATIC_DIR.exists():
+    app.mount("/", StaticFiles(directory=str(STATIC_DIR), html=True), name="static")
+else:
+    @app.get("/")
+    def root():
+        return JSONResponse(
+            {"ok": True, "note": "frontend not built — run `npm run build` in frontend/ or use the dev server"}
+        )

backend/modality.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import io
+from typing import Any
+
+import torch
+from PIL import Image
+import numpy as np
+
+
+VISION_MODEL_TYPES = {
+    "vit", "deit", "beit", "swin", "swinv2", "convnext", "convnextv2",
+    "dinov2", "dinov2_with_registers", "siglip", "clip", "mobilevit", "regnet",
+    "resnet", "efficientnet",
+}
+
+AUDIO_MODEL_TYPES = {
+    "wav2vec2", "hubert", "whisper", "audio-spectrogram-transformer", "ast",
+    "speech_to_text", "unispeech", "unispeech-sat", "data2vec-audio",
+}
+
+
+def _placeholder_image(size: int = 224) -> Image.Image:
+    """Generate a small synthetic image — a soft gradient with a circle.
+
+    Avoids bundling binary assets and keeps the Space self-contained."""
+    arr = np.zeros((size, size, 3), dtype=np.uint8)
+    yy, xx = np.mgrid[0:size, 0:size]
+    arr[..., 0] = (xx * 255 / size).astype(np.uint8)
+    arr[..., 1] = (yy * 255 / size).astype(np.uint8)
+    arr[..., 2] = 128
+    cy, cx = size // 2, size // 2
+    mask = (yy - cy) ** 2 + (xx - cx) ** 2 < (size // 4) ** 2
+    arr[mask] = [240, 240, 240]
+    return Image.fromarray(arr)
+
+
+def _placeholder_audio(seconds: float = 1.0, sr: int = 16000) -> np.ndarray:
+    t = np.linspace(0, seconds, int(seconds * sr), endpoint=False)
+    return (0.1 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)
+
+
+def detect_modality(model_type: str | None) -> str:
+    if not model_type:
+        return "text"
+    mt = model_type.lower()
+    if mt in VISION_MODEL_TYPES:
+        return "vision"
+    if mt in AUDIO_MODEL_TYPES:
+        return "audio"
+    return "text"
+
+
+def build_inputs(
+    modality: str,
+    tokenizer: Any,
+    processor: Any,
+    config: Any,
+    prompt: str | None,
+) -> tuple[dict, dict]:
+    """Returns (model_inputs, ui_payload).
+
+    ui_payload is the human-readable description of the input shown in the TokenStrip
+    (token IDs + decoded pieces, or image preview metadata)."""
+    if modality == "vision" and processor is not None:
+        img = _placeholder_image()
+        inputs = processor(images=img, return_tensors="pt")
+        ui = {
+            "kind": "image",
+            "size": list(img.size),
+            "shape": list(inputs["pixel_values"].shape),
+            "patch_size": getattr(config, "patch_size", None),
+        }
+        return dict(inputs), ui
+
+    if modality == "audio" and processor is not None:
+        sr = getattr(processor, "sampling_rate", 16000) or 16000
+        wav = _placeholder_audio(sr=sr)
+        try:
+            inputs = processor(wav, sampling_rate=sr, return_tensors="pt")
+        except TypeError:
+            inputs = processor(wav, return_tensors="pt")
+        first = next(iter(inputs.values()))
+        ui = {
+            "kind": "audio",
+            "sample_rate": sr,
+            "duration_s": float(len(wav) / sr),
+            "shape": list(first.shape),
+        }
+        return dict(inputs), ui
+
+    # text
+    text = prompt or "The quick brown fox jumps over the lazy dog."
+    enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=64)
+    ids = enc["input_ids"][0].tolist()
+    pieces = []
+    for i in ids:
+        try:
+            pieces.append(tokenizer.decode([i], skip_special_tokens=False))
+        except Exception:
+            pieces.append(f"<{i}>")
+    ui = {
+        "kind": "text",
+        "text": text,
+        "ids": ids,
+        "pieces": pieces,
+        "vocab_size": getattr(tokenizer, "vocab_size", None),
+        "shape": list(enc["input_ids"].shape),
+    }
+    return dict(enc), ui

backend/model_loader.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForMaskedLM,
+    AutoModelForImageClassification,
+    AutoModelForAudioClassification,
+    AutoProcessor,
+    AutoTokenizer,
+)
+
+from .modality import detect_modality
+
+log = logging.getLogger("model_loader")
+
+
+MAX_PARAMS_DEFAULT = 1_500_000_000
+
+
+@dataclass
+class LoadedModel:
+    model_id: str
+    model: torch.nn.Module
+    config: Any
+    tokenizer: Any
+    processor: Any
+    modality: str
+    head_kind: str  # "causal_lm" | "masked_lm" | "image_classification" | "audio_classification" | "encoder"
+    param_count: int
+
+
+class _Cache:
+    """Single-slot cache. Loading another model evicts the previous one."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._current: LoadedModel | None = None
+
+    def get(self, model_id: str) -> LoadedModel | None:
+        with self._lock:
+            if self._current and self._current.model_id == model_id:
+                return self._current
+            return None
+
+    def put(self, lm: LoadedModel) -> None:
+        with self._lock:
+            self._current = lm
+
+    def current(self) -> LoadedModel | None:
+        with self._lock:
+            return self._current
+
+
+CACHE = _Cache()
+
+
+def _pick_auto_class(config: Any, modality: str):
+    """Pick the most informative auto-class for this config so we get a real head."""
+    architectures = getattr(config, "architectures", None) or []
+    arch = architectures[0] if architectures else ""
+
+    if modality == "vision":
+        if "ImageClassification" in arch or hasattr(config, "id2label"):
+            return AutoModelForImageClassification, "image_classification"
+        return AutoModel, "encoder"
+    if modality == "audio":
+        if "AudioClassification" in arch:
+            return AutoModelForAudioClassification, "audio_classification"
+        return AutoModel, "encoder"
+
+    # text
+    if "ForCausalLM" in arch:
+        return AutoModelForCausalLM, "causal_lm"
+    if "ForMaskedLM" in arch:
+        return AutoModelForMaskedLM, "masked_lm"
+    # default — encoder backbone (BERT-like) without head
+    return AutoModel, "encoder"
+
+
+def _count_params(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+def load_model(model_id: str) -> LoadedModel:
+    cached = CACHE.get(model_id)
+    if cached is not None:
+        return cached
+
+    log.info("Loading config for %s", model_id)
+    config = AutoConfig.from_pretrained(model_id, trust_remote_code=False)
+    modality = detect_modality(getattr(config, "model_type", None))
+
+    auto_cls, head_kind = _pick_auto_class(config, modality)
+    log.info("Loading %s as %s (%s)", model_id, auto_cls.__name__, head_kind)
+
+    max_params = int(os.environ.get("MAX_PARAMS", MAX_PARAMS_DEFAULT))
+
+    # Load on CPU in float32. Keep it deterministic for the visualizer.
+    model = auto_cls.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        trust_remote_code=False,
+    )
+    n = _count_params(model)
+    if n > max_params:
+        del model
+        raise ValueError(
+            f"Model {model_id} has {n:,} params, exceeds MAX_PARAMS={max_params:,}."
+        )
+    model.eval()
+
+    tokenizer = None
+    processor = None
+    if modality == "text":
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+        except Exception:
+            tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+    else:
+        try:
+            processor = AutoProcessor.from_pretrained(model_id)
+        except Exception as e:
+            log.warning("AutoProcessor failed for %s: %s", model_id, e)
+
+    lm = LoadedModel(
+        model_id=model_id,
+        model=model,
+        config=config,
+        tokenizer=tokenizer,
+        processor=processor,
+        modality=modality,
+        head_kind=head_kind,
+        param_count=n,
+    )
+    CACHE.put(lm)
+    return lm

backend/runner.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import logging
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterator
+
+import torch
+import torch.nn as nn
+
+from .model_loader import LoadedModel
+from .modality import build_inputs
+from .snapshots import attention_snapshot, tensor_snapshot, topk_from_logits
+
+log = logging.getLogger("runner")
+
+
+@dataclass
+class RunEvent:
+    step: int
+    path: str
+    module_class: str
+    kind: str  # "module" | "attention" | "logits" | "input" | "done"
+    payload: dict
+
+
+@dataclass
+class RunResult:
+    events: list[RunEvent] = field(default_factory=list)
+
+
+def _decoder_for(loaded: LoadedModel) -> Callable[[int], str]:
+    tok = loaded.tokenizer
+    if tok is not None:
+        def decode(i: int) -> str:
+            return tok.decode([i], skip_special_tokens=False)
+        return decode
+    id2label = getattr(loaded.config, "id2label", None) or {}
+    def decode(i: int) -> str:
+        return str(id2label.get(i) or id2label.get(str(i)) or f"<id:{i}>")
+    return decode
+
+
+def _classify(cls_name: str) -> str:
+    from .tracer import classify
+    return classify(cls_name)
+
+
+def run_model(loaded: LoadedModel, prompt: str | None) -> RunResult:
+    """Run a forward pass with hooks and return the ordered list of events."""
+    model = loaded.model
+    inputs, ui = build_inputs(
+        loaded.modality, loaded.tokenizer, loaded.processor, loaded.config, prompt,
+    )
+
+    result = RunResult()
+    result.events.append(
+        RunEvent(step=0, path="<input>", module_class="Input", kind="input", payload=ui)
+    )
+
+    counter = {"i": 1}
+    handles: list[Any] = []
+    seen: set[int] = set()  # avoid double-recording the same module
+
+    def make_hook(path: str, cls_name: str):
+        def hook(module: nn.Module, inputs_, output):
+            mid = id(module)
+            if mid in seen:
+                return
+            seen.add(mid)
+            snap = tensor_snapshot(output)
+            if snap is None:
+                return
+            ev = RunEvent(
+                step=counter["i"],
+                path=path,
+                module_class=cls_name,
+                kind="module",
+                payload={"snapshot": snap, "kind": _classify(cls_name)},
+            )
+            counter["i"] += 1
+            result.events.append(ev)
+        return hook
+
+    for path, mod in model.named_modules():
+        if path == "":
+            continue
+        cls_name = type(mod).__name__
+        h = mod.register_forward_hook(make_hook(path, cls_name))
+        handles.append(h)
+
+    try:
+        forward_kwargs = dict(inputs)
+        # Only pass these flags if the model accepts them.
+        sig_params = getattr(model.forward, "__wrapped__", model.forward).__code__.co_varnames
+        if "output_attentions" in sig_params:
+            forward_kwargs["output_attentions"] = True
+        if "output_hidden_states" in sig_params:
+            forward_kwargs["output_hidden_states"] = True
+        with torch.no_grad():
+            outputs = model(**forward_kwargs)
+    finally:
+        for h in handles:
+            h.remove()
+
+    # Attention snapshots (cleaner than the raw module output).
+    attentions = getattr(outputs, "attentions", None)
+    if attentions is not None:
+        for layer_i, attn in enumerate(attentions):
+            snap = attention_snapshot(attn)
+            if snap is None:
+                continue
+            result.events.append(
+                RunEvent(
+                    step=counter["i"],
+                    path=f"<attention.layer.{layer_i}>",
+                    module_class="Attention",
+                    kind="attention",
+                    payload={"snapshot": snap, "layer": layer_i},
+                )
+            )
+            counter["i"] += 1
+
+    # Final logits / probabilities.
+    logits = getattr(outputs, "logits", None)
+    if logits is not None:
+        decode = _decoder_for(loaded)
+        last_logits = logits
+        if logits.dim() >= 3:
+            # LM heads: take the last token's logits.
+            last_logits = logits[:, -1, :]
+        topk = topk_from_logits(last_logits, decode_fn=decode, k=10)
+        snap = tensor_snapshot(logits)
+        result.events.append(
+            RunEvent(
+                step=counter["i"],
+                path="<output.logits>",
+                module_class="Logits",
+                kind="logits",
+                payload={"snapshot": snap, "top_k": topk, "head_kind": loaded.head_kind},
+            )
+        )
+        counter["i"] += 1
+
+    result.events.append(
+        RunEvent(
+            step=counter["i"],
+            path="<done>",
+            module_class="Done",
+            kind="done",
+            payload={"total_steps": counter["i"]},
+        )
+    )
+    return result
+
+
+# In-memory store of the most recent run so the WS endpoint can replay it.
+_RUN_LOCK = threading.Lock()
+_LAST_RUN: RunResult | None = None
+
+
+def store_run(result: RunResult) -> None:
+    global _LAST_RUN
+    with _RUN_LOCK:
+        _LAST_RUN = result
+
+
+def get_last_run() -> RunResult | None:
+    with _RUN_LOCK:
+        return _LAST_RUN
+
+
+def iter_events(result: RunResult) -> Iterator[RunEvent]:
+    yield from result.events

backend/snapshots.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+
+HEATMAP_MAX = 64
+SAMPLE_FLAT = 64
+TOPK = 10
+
+
+def _to_numpy(t: torch.Tensor) -> np.ndarray:
+    return t.detach().to(torch.float32).cpu().numpy()
+
+
+def _downsample_2d(arr: np.ndarray, max_side: int = HEATMAP_MAX) -> np.ndarray:
+    h, w = arr.shape
+    if h <= max_side and w <= max_side:
+        return arr
+    rh = max(1, h // max_side)
+    rw = max(1, w // max_side)
+    h2 = (h // rh) * rh
+    w2 = (w // rw) * rw
+    arr = arr[:h2, :w2]
+    return arr.reshape(h2 // rh, rh, w2 // rw, rw).mean(axis=(1, 3))
+
+
+def _to_2d(arr: np.ndarray) -> np.ndarray | None:
+    if arr.ndim == 0:
+        return None
+    if arr.ndim == 1:
+        return arr.reshape(1, -1)
+    if arr.ndim == 2:
+        return arr
+    # Collapse all leading dims into rows; last dim is columns.
+    return arr.reshape(-1, arr.shape[-1])
+
+
+def tensor_snapshot(t: Any) -> dict | None:
+    """Convert a tensor (or first tensor inside a tuple/list) into a JSON-safe summary."""
+    if isinstance(t, (tuple, list)):
+        for item in t:
+            snap = tensor_snapshot(item)
+            if snap is not None:
+                return snap
+        return None
+    if not isinstance(t, torch.Tensor):
+        return None
+    arr = _to_numpy(t)
+    flat = arr.reshape(-1)
+    if flat.size == 0:
+        return None
+    finite = flat[np.isfinite(flat)]
+    if finite.size == 0:
+        finite = np.zeros(1, dtype=np.float32)
+    stats = {
+        "mean": float(finite.mean()),
+        "std": float(finite.std()),
+        "min": float(finite.min()),
+        "max": float(finite.max()),
+        "abs_max": float(np.abs(finite).max()),
+    }
+    sample = flat[:SAMPLE_FLAT].tolist()
+    heatmap_payload: dict | None = None
+    arr2 = _to_2d(arr)
+    if arr2 is not None and arr2.size > 0:
+        ds = _downsample_2d(arr2)
+        heatmap_payload = {
+            "shape": list(ds.shape),
+            "values": ds.astype(np.float32).round(5).tolist(),
+            "vmin": float(ds.min()),
+            "vmax": float(ds.max()),
+        }
+    return {
+        "shape": list(arr.shape),
+        "dtype": str(t.dtype).replace("torch.", ""),
+        "stats": stats,
+        "sample": [round(float(x), 5) for x in sample],
+        "heatmap": heatmap_payload,
+    }
+
+
+def topk_from_logits(logits: torch.Tensor, decode_fn, k: int = TOPK) -> list[dict]:
+    """Take logits over a vocab axis and return decoded top-k probabilities."""
+    if logits.dim() == 0:
+        return []
+    flat = logits
+    while flat.dim() > 1:
+        flat = flat[0] if flat.size(0) > 0 else flat.squeeze(0)
+    probs = torch.softmax(flat.float(), dim=-1)
+    top = torch.topk(probs, k=min(k, probs.numel()))
+    out = []
+    for prob, idx in zip(top.values.tolist(), top.indices.tolist()):
+        try:
+            decoded = decode_fn(int(idx))
+        except Exception:
+            decoded = f"<id:{idx}>"
+        out.append({"id": int(idx), "token": decoded, "prob": float(prob)})
+    return out
+
+
+def attention_snapshot(attn: torch.Tensor) -> dict | None:
+    """Attention is [B, H, T, T]. Return [H, T, T] downsampled per head, plus a head-mean."""
+    if not isinstance(attn, torch.Tensor) or attn.dim() < 3:
+        return None
+    a = attn.detach().to(torch.float32).cpu()
+    if a.dim() == 4:
+        a = a[0]  # batch 0
+    if a.dim() == 3:
+        head_mean = a.mean(dim=0).numpy()
+    else:
+        head_mean = a.numpy()
+    ds = _downsample_2d(head_mean, max_side=HEATMAP_MAX)
+    return {
+        "shape": list(attn.shape),
+        "heatmap": {
+            "shape": list(ds.shape),
+            "values": ds.astype(np.float32).round(5).tolist(),
+            "vmin": float(ds.min()),
+            "vmax": float(ds.max()),
+        },
+    }

backend/tracer.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+# Coarse module-class buckets used for graph styling on the frontend.
+KIND_PATTERNS: list[tuple[str, tuple[str, ...]]] = [
+    ("Embedding", ("Embedding", "Embeddings", "PatchEmbed", "PositionalEmbedding")),
+    ("Attention", ("Attention", "SelfAttention", "MHA", "MultiHead")),
+    ("MLP", ("MLP", "FFN", "FeedForward", "Intermediate", "GLU", "MoE", "Expert")),
+    ("Norm", ("LayerNorm", "RMSNorm", "BatchNorm", "GroupNorm", "Norm")),
+    ("Linear", ("Linear",)),
+    ("Conv", ("Conv1d", "Conv2d", "Conv3d", "Conv")),
+    ("Activation", ("GELU", "ReLU", "SiLU", "Mish", "Swish", "Tanh", "Sigmoid")),
+    ("Dropout", ("Dropout",)),
+    ("Pooler", ("Pool", "Pooler")),
+    ("Block", ("Block", "Layer", "Encoder", "Decoder", "DecoderLayer", "EncoderLayer")),
+    ("Head", ("Head", "Classifier", "LMHead", "Predictor", "Projector")),
+]
+
+
+def classify(module_class: str) -> str:
+    for kind, patterns in KIND_PATTERNS:
+        for p in patterns:
+            if p in module_class:
+                return kind
+    return "Other"
+
+
+def _interesting_config(module: nn.Module) -> dict[str, Any]:
+    """Pull a few attributes off the module that are visualizer-friendly."""
+    out: dict[str, Any] = {}
+    for k in (
+        "num_heads",
+        "num_attention_heads",
+        "head_dim",
+        "hidden_size",
+        "intermediate_size",
+        "embed_dim",
+        "in_features",
+        "out_features",
+        "in_channels",
+        "out_channels",
+        "kernel_size",
+        "stride",
+        "padding",
+        "eps",
+        "normalized_shape",
+        "num_embeddings",
+        "embedding_dim",
+        "num_experts",
+        "top_k",
+    ):
+        v = getattr(module, k, None)
+        if v is None:
+            continue
+        if isinstance(v, (int, float, str, bool)):
+            out[k] = v
+        elif isinstance(v, (list, tuple)) and all(isinstance(x, (int, float)) for x in v):
+            out[k] = list(v)
+    return out
+
+
+def build_graph(model: nn.Module) -> dict:
+    """Walk the module tree and return a JSON graph.
+
+    We build edges from the parent → each child in declaration order. This
+    captures the module hierarchy precisely; runtime dataflow is overlaid
+    later by the runner using the order of forward-hook firings.
+    """
+    nodes: list[dict] = []
+    edges: list[dict] = []
+
+    # Collect each module's parameter count cheaply.
+    param_counts: dict[str, int] = {}
+    for path, mod in model.named_modules():
+        param_counts[path] = sum(p.numel() for p in mod.parameters(recurse=False))
+
+    # Aggregate: every node also reports the *recursive* parameter count.
+    recursive: dict[str, int] = {}
+
+    def _walk(prefix: str, mod: nn.Module) -> int:
+        own = param_counts.get(prefix, 0)
+        total = own
+        for name, child in mod.named_children():
+            child_path = f"{prefix}.{name}" if prefix else name
+            total += _walk(child_path, child)
+        recursive[prefix] = total
+        return total
+
+    _walk("", model)
+
+    # Emit nodes + parent→child edges.
+    for path, mod in model.named_modules():
+        cls = type(mod).__name__
+        node_id = path or "<root>"
+        children = [
+            f"{path}.{name}" if path else name for name, _ in mod.named_children()
+        ]
+        nodes.append(
+            {
+                "id": node_id,
+                "path": path,
+                "module_class": cls,
+                "kind": classify(cls),
+                "params": recursive.get(path, 0),
+                "params_own": param_counts.get(path, 0),
+                "config": _interesting_config(mod),
+                "children": [c if c else "<root>" for c in children],
+                "depth": 0 if not path else path.count(".") + 1,
+                "is_leaf": len(children) == 0,
+            }
+        )
+        for c in children:
+            edges.append(
+                {
+                    "from": node_id,
+                    "to": c if c else "<root>",
+                    "kind": "tree",
+                }
+            )
+
+    return {
+        "root": "<root>",
+        "nodes": nodes,
+        "edges": edges,
+    }

frontend/index.html

@@ -0,0 +1,12 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Transformer X-Ray</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>

frontend/package.json

@@ -0,0 +1,27 @@
+{
+  "name": "transformer-xray",
+  "private": true,
+  "version": "0.1.0",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "tsc -b --noEmit && vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "@dagrejs/dagre": "^1.1.4",
+    "gif.js": "^0.2.0",
+    "html-to-image": "^1.11.13",
+    "react": "^18.3.1",
+    "react-dom": "^18.3.1",
+    "reactflow": "^11.11.4",
+    "zustand": "^4.5.5"
+  },
+  "devDependencies": {
+    "@types/react": "^18.3.12",
+    "@types/react-dom": "^18.3.1",
+    "@vitejs/plugin-react": "^4.3.4",
+    "typescript": "^5.6.3",
+    "vite": "^5.4.11"
+  }
+}

frontend/src/App.tsx

@@ -0,0 +1,36 @@
+import Graph from "./components/Graph";
+import Header from "./components/Header";
+import NodeCard from "./components/NodeCard";
+import PlayBar from "./components/PlayBar";
+import ProbsBar from "./components/ProbsBar";
+import TokenStrip from "./components/TokenStrip";
+import GifExport from "./components/GifExport";
+import { useStore } from "./store";
+
+export default function App() {
+  const error = useStore((s) => s.error);
+  return (
+    <div style={{ display: "grid", gridTemplateRows: "auto auto 1fr auto", height: "100vh", background: "#0b1220", color: "#e2e8f0" }}>
+      <Header />
+      <TokenStrip />
+      <main style={{ display: "grid", gridTemplateColumns: "1fr 360px", overflow: "hidden" }}>
+        <Graph />
+        <aside style={{ borderLeft: "1px solid #1f2937", background: "#0f172a", overflow: "auto", display: "flex", flexDirection: "column" }}>
+          <div style={{ borderBottom: "1px solid #1f2937" }}>
+            <NodeCard />
+          </div>
+          <div style={{ borderBottom: "1px solid #1f2937" }}>
+            <ProbsBar />
+          </div>
+          <GifExport />
+        </aside>
+      </main>
+      <PlayBar />
+      {error && (
+        <div style={{ position: "fixed", bottom: 80, right: 16, background: "#7f1d1d", color: "#fee2e2", padding: "8px 12px", borderRadius: 6, fontSize: 12, maxWidth: 360 }}>
+          {error}
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/components/GifExport.tsx

@@ -0,0 +1,99 @@
+import { useState } from "react";
+// @ts-ignore — gif.js has no types
+import GIF from "gif.js";
+import { toPng } from "html-to-image";
+import { useStore } from "../store";
+
+export default function GifExport() {
+  const events = useStore((s) => s.events);
+  const setCurrentStep = useStore((s) => s.setCurrentStep);
+  const [busy, setBusy] = useState(false);
+  const [progress, setProgress] = useState(0);
+  const [fps, setFps] = useState(8);
+  const [maxFrames, setMaxFrames] = useState(60);
+
+  async function record() {
+    if (busy || events.length === 0) return;
+    setBusy(true);
+    setProgress(0);
+    try {
+      const target = document.getElementById("graph-canvas");
+      if (!target) throw new Error("graph not found");
+
+      const total = Math.min(maxFrames, events.length);
+      const stride = Math.max(1, Math.floor(events.length / total));
+
+      // gif.worker.js is bundled with the package; resolve at runtime so Vite serves it.
+      const workerUrl = new URL(
+        "gif.js/dist/gif.worker.js",
+        import.meta.url,
+      ).toString();
+      const gif = new GIF({
+        workers: 2,
+        quality: 8,
+        workerScript: workerUrl,
+        width: target.clientWidth,
+        height: target.clientHeight,
+      });
+
+      for (let i = 0; i < total; i++) {
+        const step = Math.min(events.length - 1, i * stride);
+        setCurrentStep(step);
+        // Allow React to repaint before snapshot.
+        await new Promise((r) => requestAnimationFrame(() => requestAnimationFrame(r)));
+        const dataUrl = await toPng(target, { cacheBust: true, pixelRatio: 1 });
+        const img = await new Promise<HTMLImageElement>((resolve, reject) => {
+          const im = new Image();
+          im.onload = () => resolve(im);
+          im.onerror = reject;
+          im.src = dataUrl;
+        });
+        gif.addFrame(img, { delay: Math.round(1000 / fps) });
+        setProgress(Math.round(((i + 1) / total) * 50));
+      }
+
+      gif.on("progress", (p: number) => setProgress(50 + Math.round(p * 50)));
+      const blob: Blob = await new Promise((resolve) => {
+        gif.on("finished", (b: Blob) => resolve(b));
+        gif.render();
+      });
+
+      const url = URL.createObjectURL(blob);
+      const a = document.createElement("a");
+      a.href = url;
+      a.download = `transformer-xray-${Date.now()}.gif`;
+      a.click();
+      URL.revokeObjectURL(url);
+    } catch (e) {
+      console.error(e);
+      alert(`gif export failed: ${(e as Error).message}`);
+    } finally {
+      setBusy(false);
+      setProgress(0);
+    }
+  }
+
+  return (
+    <div style={{ padding: 12, borderTop: "1px solid #1f2937" }}>
+      <div style={{ fontSize: 11, color: "#94a3b8", marginBottom: 6 }}>GIF export</div>
+      <div style={{ display: "flex", gap: 8, alignItems: "center", marginBottom: 6, fontSize: 11 }}>
+        <label style={{ color: "#94a3b8" }}>
+          fps
+          <input type="number" min={2} max={24} value={fps} onChange={(e) => setFps(Math.max(2, Math.min(24, Number(e.target.value))))} style={{ width: 50, marginLeft: 6 }} />
+        </label>
+        <label style={{ color: "#94a3b8" }}>
+          frames
+          <input type="number" min={10} max={200} value={maxFrames} onChange={(e) => setMaxFrames(Math.max(10, Math.min(200, Number(e.target.value))))} style={{ width: 60, marginLeft: 6 }} />
+        </label>
+      </div>
+      <button className="btn btn-primary" disabled={busy || events.length === 0} onClick={record}>
+        {busy ? `recording… ${progress}%` : "record GIF"}
+      </button>
+      {busy && (
+        <div style={{ marginTop: 6, height: 4, background: "#0b1220", borderRadius: 2, overflow: "hidden" }}>
+          <div style={{ width: `${progress}%`, height: "100%", background: "#38bdf8", transition: "width 120ms" }} />
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/components/Graph.tsx

@@ -0,0 +1,45 @@
+import { useMemo } from "react";
+import ReactFlow, { Background, Controls, MiniMap } from "reactflow";
+import { useStore } from "../store";
+import { layoutGraph } from "../lib/layout";
+import ModuleNode from "./ModuleNode";
+
+const nodeTypes = { module: ModuleNode };
+
+export default function Graph() {
+  const graph = useStore((s) => s.graph);
+  const granularity = useStore((s) => s.granularity);
+  const setSelectedPath = useStore((s) => s.setSelectedPath);
+
+  const { nodes, edges } = useMemo(() => {
+    if (!graph) return { nodes: [], edges: [] };
+    return layoutGraph(graph, { granularity });
+  }, [graph, granularity]);
+
+  if (!graph) {
+    return (
+      <div style={{ display: "grid", placeItems: "center", height: "100%", color: "#94a3b8" }}>
+        Load a model to see its graph.
+      </div>
+    );
+  }
+
+  return (
+    <div id="graph-canvas" style={{ width: "100%", height: "100%", background: "#0b1220" }}>
+      <ReactFlow
+        nodes={nodes}
+        edges={edges}
+        nodeTypes={nodeTypes}
+        fitView
+        minZoom={0.05}
+        maxZoom={2}
+        onPaneClick={() => setSelectedPath(null)}
+        proOptions={{ hideAttribution: true }}
+      >
+        <Background color="#1f2937" gap={24} />
+        <Controls showInteractive={false} />
+        <MiniMap nodeColor={() => "#475569"} maskColor="rgba(11,18,32,0.7)" pannable zoomable />
+      </ReactFlow>
+    </div>
+  );
+}

frontend/src/components/Header.tsx

@@ -0,0 +1,104 @@
+import { useState } from "react";
+import { useStore } from "../store";
+import { fetchGraph, loadModel } from "../lib/api";
+
+const PRESETS = [
+  "prajjwal1/bert-tiny",
+  "bert-base-uncased",
+  "distilgpt2",
+  "google/vit-base-patch16-224",
+  "openai-community/gpt2",
+];
+
+function fmt(n: number): string {
+  if (n >= 1e9) return (n / 1e9).toFixed(2) + "B";
+  if (n >= 1e6) return (n / 1e6).toFixed(2) + "M";
+  if (n >= 1e3) return (n / 1e3).toFixed(1) + "K";
+  return String(n);
+}
+
+export default function Header() {
+  const modelId = useStore((s) => s.modelId);
+  const setModelId = useStore((s) => s.setModelId);
+  const setLoadInfo = useStore((s) => s.setLoadInfo);
+  const setGraph = useStore((s) => s.setGraph);
+  const setLoading = useStore((s) => s.setLoading);
+  const setError = useStore((s) => s.setError);
+  const resetRun = useStore((s) => s.resetRun);
+  const loading = useStore((s) => s.loading);
+  const loadInfo = useStore((s) => s.loadInfo);
+  const granularity = useStore((s) => s.granularity);
+  const setGranularity = useStore((s) => s.setGranularity);
+  const [pending, setPending] = useState(modelId);
+
+  async function go(id: string) {
+    setLoading(true);
+    setError(null);
+    setLoadInfo(null);
+    setGraph(null);
+    resetRun();
+    try {
+      const info = await loadModel(id);
+      setLoadInfo(info);
+      const g = await fetchGraph();
+      setGraph(g);
+      setModelId(id);
+    } catch (e: any) {
+      setError(e?.message ?? String(e));
+    } finally {
+      setLoading(false);
+    }
+  }
+
+  return (
+    <header style={{ padding: "10px 14px", borderBottom: "1px solid #1f2937", background: "#0f172a", display: "flex", flexDirection: "column", gap: 8 }}>
+      <div style={{ display: "flex", alignItems: "center", gap: 8, flexWrap: "wrap" }}>
+        <span style={{ fontWeight: 800, color: "#f8fafc", letterSpacing: 0.3 }}>🔬 Transformer X-Ray</span>
+        <input
+          value={pending}
+          onChange={(e) => setPending(e.target.value)}
+          placeholder="huggingface-id e.g. bert-base-uncased"
+          onKeyDown={(e) => { if (e.key === "Enter") go(pending); }}
+          style={{
+            flex: 1,
+            minWidth: 280,
+            background: "#0b1220",
+            color: "#e2e8f0",
+            border: "1px solid #334155",
+            borderRadius: 6,
+            padding: "6px 10px",
+            fontFamily: "ui-monospace, monospace",
+            fontSize: 12,
+          }}
+        />
+        <button className="btn btn-primary" onClick={() => go(pending)} disabled={loading}>
+          {loading ? "loading…" : "load"}
+        </button>
+        <div style={{ display: "flex", gap: 4, marginLeft: 6 }}>
+          <span style={{ fontSize: 10, color: "#94a3b8", alignSelf: "center" }}>granularity</span>
+          {[0, 1, 2].map((g) => (
+            <button key={g} className={g === granularity ? "tab tab-active" : "tab"} onClick={() => setGranularity(g)}>
+              {["top", "blocks", "leaves"][g]}
+            </button>
+          ))}
+        </div>
+      </div>
+      <div style={{ display: "flex", gap: 4, flexWrap: "wrap", fontSize: 10 }}>
+        {PRESETS.map((p) => (
+          <button key={p} className="tab" onClick={() => { setPending(p); go(p); }}>{p}</button>
+        ))}
+        {loadInfo && (
+          <div style={{ marginLeft: "auto", color: "#94a3b8", fontFamily: "ui-monospace, monospace" }}>
+            {loadInfo.config.model_type} ·
+            {" "}{fmt(loadInfo.param_count)} params ·
+            {" "}{loadInfo.modality} ·
+            {" "}head <span style={{ color: "#fb923c" }}>{loadInfo.head_kind}</span>
+            {loadInfo.config.hidden_size && <> · hidden {loadInfo.config.hidden_size}</>}
+            {loadInfo.config.num_hidden_layers && <> · L={loadInfo.config.num_hidden_layers}</>}
+            {loadInfo.config.num_attention_heads && <> · H={loadInfo.config.num_attention_heads}</>}
+          </div>
+        )}
+      </div>
+    </header>
+  );
+}

frontend/src/components/ModuleNode.tsx

@@ -0,0 +1,63 @@
+import { Handle, Position, type NodeProps } from "reactflow";
+import { useStore } from "../store";
+import { colorFor } from "../lib/colors";
+import type { GraphNode } from "../types";
+
+function fmtParams(n: number): string {
+  if (n >= 1e9) return (n / 1e9).toFixed(2) + "B";
+  if (n >= 1e6) return (n / 1e6).toFixed(2) + "M";
+  if (n >= 1e3) return (n / 1e3).toFixed(1) + "K";
+  return String(n);
+}
+
+export default function ModuleNode({ data, id }: NodeProps<GraphNode>) {
+  const events = useStore((s) => s.events);
+  const currentStep = useStore((s) => s.currentStep);
+  const snapshotsByPath = useStore((s) => s.snapshotsByPath);
+  const selectedPath = useStore((s) => s.selectedPath);
+  const setSelectedPath = useStore((s) => s.setSelectedPath);
+
+  const c = colorFor(data.kind);
+  const ev = events[currentStep];
+  const isActive = ev && ev.path === data.path && ev.kind === "module";
+  const hasFlowed = !!snapshotsByPath[data.path];
+  const isSelected = selectedPath === id;
+
+  const snap = snapshotsByPath[data.path];
+  const shapeStr = snap ? `[${snap.shape.join(", ")}]` : "";
+
+  return (
+    <div
+      onClick={(e) => {
+        e.stopPropagation();
+        setSelectedPath(id);
+      }}
+      style={{
+        width: 220,
+        padding: "8px 10px",
+        borderRadius: 8,
+        background: c.bg,
+        border: `2px solid ${isSelected ? "#fff" : isActive ? c.accent : c.border}`,
+        color: c.text,
+        boxShadow: isActive ? `0 0 24px ${c.accent}` : hasFlowed ? `0 0 0 1px ${c.accent}` : "none",
+        opacity: hasFlowed || !ev ? 1 : 0.55,
+        transition: "box-shadow 120ms, opacity 120ms, border-color 120ms",
+        fontSize: 11,
+        lineHeight: 1.35,
+        fontFamily: "ui-monospace, Menlo, monospace",
+      }}
+    >
+      <Handle type="target" position={Position.Top} style={{ background: c.accent }} />
+      <div style={{ display: "flex", justifyContent: "space-between", gap: 6 }}>
+        <span style={{ fontWeight: 700, color: c.accent }}>{data.kind}</span>
+        <span style={{ opacity: 0.7 }}>{fmtParams(data.params)}</span>
+      </div>
+      <div style={{ fontWeight: 600, fontSize: 12, marginTop: 2 }}>{data.module_class}</div>
+      <div style={{ opacity: 0.7, fontSize: 10, marginTop: 2 }}>{data.path || "<root>"}</div>
+      {shapeStr && (
+        <div style={{ marginTop: 4, fontSize: 10, color: c.accent, fontWeight: 600 }}>{shapeStr}</div>
+      )}
+      <Handle type="source" position={Position.Bottom} style={{ background: c.accent }} />
+    </div>
+  );
+}

frontend/src/components/NodeCard.tsx

@@ -0,0 +1,64 @@
+import { useStore } from "../store";
+import TensorView from "./TensorView";
+
+function fmtParams(n: number): string {
+  if (n >= 1e9) return (n / 1e9).toFixed(2) + "B";
+  if (n >= 1e6) return (n / 1e6).toFixed(2) + "M";
+  if (n >= 1e3) return (n / 1e3).toFixed(1) + "K";
+  return String(n);
+}
+
+export default function NodeCard() {
+  const graph = useStore((s) => s.graph);
+  const selectedPath = useStore((s) => s.selectedPath);
+  const snapshotsByPath = useStore((s) => s.snapshotsByPath);
+
+  if (!graph) return null;
+  const node = selectedPath ? graph.nodes.find((n) => n.id === selectedPath) : null;
+  const snap = node ? snapshotsByPath[node.path] : undefined;
+
+  if (!node) {
+    return (
+      <div style={{ padding: 12, color: "#64748b", fontSize: 12 }}>
+        Click a node to see its config and the tensor that flowed through it.
+      </div>
+    );
+  }
+
+  return (
+    <div style={{ padding: 12, fontSize: 12, color: "#e2e8f0" }}>
+      <div style={{ fontSize: 13, fontWeight: 700, color: "#f8fafc" }}>{node.module_class}</div>
+      <div style={{ fontSize: 10, opacity: 0.6, fontFamily: "ui-monospace, monospace", marginTop: 2, wordBreak: "break-all" }}>
+        {node.path || "<root>"}
+      </div>
+      <div style={{ display: "flex", gap: 8, marginTop: 8, fontSize: 11 }}>
+        <span className="pill">{node.kind}</span>
+        <span className="pill">{fmtParams(node.params)} params</span>
+        {node.is_leaf && <span className="pill">leaf</span>}
+      </div>
+
+      {Object.keys(node.config).length > 0 && (
+        <div style={{ marginTop: 10 }}>
+          <div style={{ color: "#94a3b8", fontSize: 11, marginBottom: 4 }}>module config</div>
+          <div style={{ display: "grid", gridTemplateColumns: "auto 1fr", gap: "2px 8px", fontFamily: "ui-monospace, monospace", fontSize: 11 }}>
+            {Object.entries(node.config).map(([k, v]) => (
+              <div key={k} style={{ display: "contents" }}>
+                <span style={{ color: "#94a3b8" }}>{k}</span>
+                <span style={{ color: "#e2e8f0" }}>{JSON.stringify(v)}</span>
+              </div>
+            ))}
+          </div>
+        </div>
+      )}
+
+      <div style={{ marginTop: 12, color: "#94a3b8", fontSize: 11 }}>output tensor</div>
+      {snap ? (
+        <TensorView snap={snap} />
+      ) : (
+        <div style={{ marginTop: 6, color: "#64748b", fontSize: 11 }}>
+          run the model to see this module's output
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/components/PlayBar.tsx

@@ -0,0 +1,118 @@
+import { useStore } from "../store";
+import { usePlayback, useRunController } from "../hooks/useRun";
+import { colorFor } from "../lib/colors";
+
+const SPEEDS = [0.5, 1, 2, 4];
+
+export default function PlayBar() {
+  usePlayback();
+  const { start } = useRunController();
+  const events = useStore((s) => s.events);
+  const currentStep = useStore((s) => s.currentStep);
+  const playing = useStore((s) => s.playing);
+  const speed = useStore((s) => s.speed);
+  const setPlaying = useStore((s) => s.setPlaying);
+  const setSpeed = useStore((s) => s.setSpeed);
+  const setCurrentStep = useStore((s) => s.setCurrentStep);
+  const running = useStore((s) => s.running);
+  const loaded = useStore((s) => s.loadInfo);
+  const promptDefault = loaded?.modality === "text" ? "Paris is the capital of [MASK]." : "";
+
+  const ev = events[currentStep];
+  const accent = ev?.kind === "module"
+    ? colorFor(ev.payload?.kind ?? "Other").accent
+    : ev?.kind === "attention"
+      ? "#a78bfa"
+      : ev?.kind === "logits"
+        ? "#fb923c"
+        : "#38bdf8";
+
+  // Build milestone markers — first appearance of each kind + first/middle/last "Block" snapshot.
+  const milestones: { step: number; label: string }[] = [];
+  if (events.length) {
+    const seen = new Set<string>();
+    events.forEach((e, i) => {
+      if (e.kind === "input" && !seen.has("input")) { milestones.push({ step: i, label: "Tokenize" }); seen.add("input"); }
+      if (e.kind === "module" && e.payload?.kind === "Embedding" && !seen.has("emb")) { milestones.push({ step: i, label: "Embed" }); seen.add("emb"); }
+      if (e.kind === "module" && e.payload?.kind === "Block" && !seen.has("block0")) { milestones.push({ step: i, label: "Block 0" }); seen.add("block0"); }
+      if (e.kind === "module" && e.payload?.kind === "Head" && !seen.has("head")) { milestones.push({ step: i, label: "Head" }); seen.add("head"); }
+      if (e.kind === "logits" && !seen.has("logits")) { milestones.push({ step: i, label: "Logits → Softmax" }); seen.add("logits"); }
+    });
+  }
+
+  return (
+    <div style={{ padding: "10px 14px", borderTop: "1px solid #1f2937", background: "#0f172a" }}>
+      <div style={{ display: "flex", alignItems: "center", gap: 10, flexWrap: "wrap" }}>
+        <button
+          className="btn btn-primary"
+          disabled={!loaded || running}
+          onClick={() => start(promptDefault || null)}
+        >
+          {running ? "running…" : "▶ Play"}
+        </button>
+        <button
+          className="btn"
+          disabled={events.length === 0}
+          onClick={() => setPlaying(!playing)}
+        >
+          {playing ? "❚❚ pause" : "▶ resume"}
+        </button>
+        <div style={{ display: "flex", gap: 2 }}>
+          {SPEEDS.map((s) => (
+            <button key={s} className={s === speed ? "tab tab-active" : "tab"} onClick={() => setSpeed(s)}>
+              {s}×
+            </button>
+          ))}
+        </div>
+        <div style={{ flex: 1, minWidth: 200, position: "relative" }}>
+          <input
+            type="range"
+            min={0}
+            max={Math.max(0, events.length - 1)}
+            value={currentStep}
+            onChange={(e) => setCurrentStep(Number(e.target.value))}
+            style={{ width: "100%", accentColor: accent }}
+          />
+          <div style={{ position: "relative", height: 14, marginTop: -4 }}>
+            {milestones.map((m) => (
+              <button
+                key={m.label}
+                onClick={() => setCurrentStep(m.step)}
+                title={m.label}
+                style={{
+                  position: "absolute",
+                  left: `${(m.step / Math.max(1, events.length - 1)) * 100}%`,
+                  transform: "translateX(-50%)",
+                  fontSize: 9,
+                  padding: "1px 6px",
+                  borderRadius: 8,
+                  background: "#1e293b",
+                  color: "#cbd5e1",
+                  border: "1px solid #334155",
+                  cursor: "pointer",
+                  whiteSpace: "nowrap",
+                }}
+              >
+                {m.label}
+              </button>
+            ))}
+          </div>
+        </div>
+        <div style={{ fontFamily: "ui-monospace, monospace", fontSize: 11, color: "#94a3b8", minWidth: 70, textAlign: "right" }}>
+          {currentStep} / {Math.max(0, events.length - 1)}
+        </div>
+      </div>
+      <div style={{ fontSize: 11, color: "#94a3b8", marginTop: 6, fontFamily: "ui-monospace, monospace" }}>
+        {ev ? (
+          <>
+            <span style={{ color: accent, fontWeight: 600 }}>[{ev.kind}]</span>{" "}
+            <span style={{ color: "#e2e8f0" }}>{ev.module_class}</span>{" "}
+            <span style={{ opacity: 0.7 }}>{ev.path}</span>
+          </>
+        ) : (
+          "no run yet — press Play to capture a forward pass"
+        )}
+      </div>
+    </div>
+  );
+}

frontend/src/components/ProbsBar.tsx

@@ -0,0 +1,55 @@
+import { useStore } from "../store";
+
+export default function ProbsBar() {
+  const topK = useStore((s) => s.topK);
+  const loadInfo = useStore((s) => s.loadInfo);
+
+  if (!topK || topK.length === 0) {
+    return (
+      <div style={{ padding: 12, color: "#64748b", fontSize: 11 }}>
+        top-k probabilities will appear here after the run completes.
+      </div>
+    );
+  }
+
+  const max = Math.max(...topK.map((t) => t.prob));
+  return (
+    <div style={{ padding: 12 }}>
+      <div style={{ fontSize: 11, color: "#94a3b8", marginBottom: 6 }}>
+        head: <span style={{ color: "#fb923c" }}>{loadInfo?.head_kind}</span>
+      </div>
+      <div style={{ display: "flex", flexDirection: "column", gap: 4 }}>
+        {topK.map((t, i) => (
+          <div key={i} style={{ display: "grid", gridTemplateColumns: "60px 1fr 60px", alignItems: "center", gap: 8 }}>
+            <span
+              style={{
+                fontFamily: "ui-monospace, monospace",
+                fontSize: 11,
+                color: "#e2e8f0",
+                whiteSpace: "nowrap",
+                overflow: "hidden",
+                textOverflow: "ellipsis",
+              }}
+              title={t.token}
+            >
+              {t.token}
+            </span>
+            <div style={{ background: "#0b1220", borderRadius: 3, height: 14, overflow: "hidden", border: "1px solid #1f2937" }}>
+              <div
+                style={{
+                  width: `${(t.prob / max) * 100}%`,
+                  height: "100%",
+                  background: i === 0 ? "#fb923c" : "#60a5fa",
+                  transition: "width 240ms",
+                }}
+              />
+            </div>
+            <span style={{ fontFamily: "ui-monospace, monospace", fontSize: 10, color: "#94a3b8", textAlign: "right" }}>
+              {(t.prob * 100).toFixed(2)}%
+            </span>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+}

frontend/src/components/TensorView.tsx

@@ -0,0 +1,122 @@
+import { useEffect, useRef, useState } from "react";
+import { viridis } from "../lib/colors";
+import type { Heatmap, Snapshot } from "../types";
+
+function HeatmapCanvas({ hm }: { hm: Heatmap }) {
+  const ref = useRef<HTMLCanvasElement>(null);
+  useEffect(() => {
+    const canvas = ref.current;
+    if (!canvas) return;
+    const [h, w] = hm.shape;
+    canvas.width = w;
+    canvas.height = h;
+    const ctx = canvas.getContext("2d");
+    if (!ctx) return;
+    const img = ctx.createImageData(w, h);
+    const span = hm.vmax - hm.vmin || 1;
+    for (let i = 0; i < h; i++) {
+      for (let j = 0; j < w; j++) {
+        const v = hm.values[i][j];
+        const t = (v - hm.vmin) / span;
+        const c = viridis(t);
+        const m = c.match(/\d+/g)!;
+        const idx = (i * w + j) * 4;
+        img.data[idx] = +m[0];
+        img.data[idx + 1] = +m[1];
+        img.data[idx + 2] = +m[2];
+        img.data[idx + 3] = 255;
+      }
+    }
+    ctx.putImageData(img, 0, 0);
+  }, [hm]);
+  return (
+    <canvas
+      ref={ref}
+      style={{
+        width: "100%",
+        maxWidth: 320,
+        imageRendering: "pixelated",
+        borderRadius: 4,
+        border: "1px solid #1f2937",
+        background: "#0b1220",
+      }}
+    />
+  );
+}
+
+export default function TensorView({ snap }: { snap: Snapshot }) {
+  const [mode, setMode] = useState<"heatmap" | "numbers" | "histogram">("heatmap");
+
+  const fmt = (x: number) => x.toFixed(Math.abs(x) < 0.01 ? 4 : 3);
+  const buckets = (() => {
+    const vals = snap.sample;
+    if (vals.length === 0) return [];
+    const min = Math.min(...vals);
+    const max = Math.max(...vals);
+    const bins = 20;
+    const span = max - min || 1;
+    const counts = new Array(bins).fill(0);
+    for (const v of vals) {
+      const b = Math.min(bins - 1, Math.floor(((v - min) / span) * bins));
+      counts[b]++;
+    }
+    const peak = Math.max(...counts) || 1;
+    return counts.map((c, i) => ({ pct: c / peak, label: min + (i / bins) * span }));
+  })();
+
+  return (
+    <div style={{ marginTop: 8 }}>
+      <div style={{ display: "flex", gap: 4, marginBottom: 6 }}>
+        {(["heatmap", "numbers", "histogram"] as const).map((m) => (
+          <button
+            key={m}
+            onClick={() => setMode(m)}
+            className={mode === m ? "tab tab-active" : "tab"}
+          >
+            {m}
+          </button>
+        ))}
+      </div>
+      <div style={{ fontSize: 11, color: "#94a3b8", marginBottom: 6, fontFamily: "ui-monospace, monospace" }}>
+        shape [{snap.shape.join(", ")}] · {snap.dtype} · μ={fmt(snap.stats.mean)} σ={fmt(snap.stats.std)} ·
+        min={fmt(snap.stats.min)} max={fmt(snap.stats.max)}
+      </div>
+      {mode === "heatmap" && (
+        snap.heatmap ? <HeatmapCanvas hm={snap.heatmap} /> : <div style={{ color: "#64748b" }}>no 2D projection</div>
+      )}
+      {mode === "numbers" && (
+        <div
+          style={{
+            display: "grid",
+            gridTemplateColumns: "repeat(8, 1fr)",
+            gap: 2,
+            fontFamily: "ui-monospace, monospace",
+            fontSize: 10,
+          }}
+        >
+          {snap.sample.map((v, i) => (
+            <span key={i} style={{ padding: "2px 4px", background: "#0b1220", borderRadius: 2, textAlign: "right", color: "#cbd5e1" }}>
+              {fmt(v)}
+            </span>
+          ))}
+        </div>
+      )}
+      {mode === "histogram" && (
+        <div style={{ display: "flex", alignItems: "flex-end", gap: 2, height: 80 }}>
+          {buckets.map((b, i) => (
+            <div
+              key={i}
+              title={fmt(b.label)}
+              style={{
+                flex: 1,
+                background: viridis(0.3 + b.pct * 0.5),
+                height: `${Math.max(2, b.pct * 100)}%`,
+                borderRadius: 1,
+              }}
+            />
+          ))}
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/components/TokenStrip.tsx

@@ -0,0 +1,68 @@
+import { useStore } from "../store";
+
+export default function TokenStrip() {
+  const inputUI = useStore((s) => s.inputUI);
+  const loadInfo = useStore((s) => s.loadInfo);
+
+  if (!inputUI) {
+    return (
+      <div style={{ padding: "8px 14px", borderBottom: "1px solid #1f2937", background: "#0f172a", color: "#64748b", fontSize: 11 }}>
+        input will appear here once you press Play
+      </div>
+    );
+  }
+
+  if (inputUI.kind === "image") {
+    return (
+      <div style={{ padding: "8px 14px", borderBottom: "1px solid #1f2937", background: "#0f172a", color: "#cbd5e1", fontSize: 11, fontFamily: "ui-monospace, monospace", display: "flex", gap: 16 }}>
+        <span><span style={{ color: "#94a3b8" }}>image</span> {inputUI.size?.[0]}×{inputUI.size?.[1]}</span>
+        <span><span style={{ color: "#94a3b8" }}>pixel_values</span> [{inputUI.shape.join(", ")}]</span>
+        {inputUI.patch_size && <span><span style={{ color: "#94a3b8" }}>patch</span> {inputUI.patch_size}</span>}
+      </div>
+    );
+  }
+  if (inputUI.kind === "audio") {
+    return (
+      <div style={{ padding: "8px 14px", borderBottom: "1px solid #1f2937", background: "#0f172a", color: "#cbd5e1", fontSize: 11, fontFamily: "ui-monospace, monospace", display: "flex", gap: 16 }}>
+        <span><span style={{ color: "#94a3b8" }}>audio</span> {inputUI.duration_s?.toFixed(2)}s @ {inputUI.sample_rate}Hz</span>
+        <span><span style={{ color: "#94a3b8" }}>shape</span> [{inputUI.shape.join(", ")}]</span>
+      </div>
+    );
+  }
+
+  const ids: number[] = inputUI.ids ?? [];
+  const pieces: string[] = inputUI.pieces ?? [];
+  const vocab = inputUI.vocab_size ?? loadInfo?.tokenizer?.vocab_size;
+
+  return (
+    <div style={{ padding: "8px 14px", borderBottom: "1px solid #1f2937", background: "#0f172a", overflowX: "auto" }}>
+      <div style={{ display: "flex", gap: 6, alignItems: "center" }}>
+        <div style={{ fontSize: 10, color: "#94a3b8", whiteSpace: "nowrap" }}>
+          tokens [{ids.length}] · vocab {vocab?.toLocaleString()}
+        </div>
+        <div style={{ display: "flex", gap: 4 }}>
+          {ids.map((id, i) => (
+            <div
+              key={i}
+              style={{
+                background: "#1e293b",
+                border: "1px solid #334155",
+                borderRadius: 6,
+                padding: "3px 6px",
+                fontSize: 11,
+                fontFamily: "ui-monospace, monospace",
+                color: "#e2e8f0",
+                whiteSpace: "nowrap",
+                lineHeight: 1.2,
+              }}
+              title={`id ${id}`}
+            >
+              <div style={{ color: "#38bdf8", fontSize: 9 }}>{id}</div>
+              <div>{pieces[i]?.replace(/^▁/, "_") ?? ""}</div>
+            </div>
+          ))}
+        </div>
+      </div>
+    </div>
+  );
+}

frontend/src/hooks/useRun.ts

@@ -0,0 +1,71 @@
+import { useEffect, useRef } from "react";
+import { useStore } from "../store";
+import { runModel, streamRun } from "../lib/api";
+
+export function useRunController() {
+  const wsRef = useRef<WebSocket | null>(null);
+  const setRunning = useStore((s) => s.setRunning);
+  const resetRun = useStore((s) => s.resetRun);
+  const appendEvent = useStore((s) => s.appendEvent);
+  const setError = useStore((s) => s.setError);
+
+  async function start(prompt: string | null) {
+    resetRun();
+    setRunning(true);
+    setError(null);
+    try {
+      const { run_id } = await runModel(prompt);
+      wsRef.current = streamRun(
+        run_id,
+        (ev) => appendEvent(ev),
+        () => setRunning(false),
+        (e) => {
+          console.error("ws error", e);
+          setError("websocket error");
+          setRunning(false);
+        },
+      );
+    } catch (e: any) {
+      setError(e?.message ?? String(e));
+      setRunning(false);
+    }
+  }
+
+  useEffect(() => {
+    return () => {
+      wsRef.current?.close();
+    };
+  }, []);
+
+  return { start };
+}
+
+/** Drives playback: advances `currentStep` over time when `playing` is true. */
+export function usePlayback() {
+  const playing = useStore((s) => s.playing);
+  const speed = useStore((s) => s.speed);
+  const events = useStore((s) => s.events);
+  const setCurrentStep = useStore((s) => s.setCurrentStep);
+  const setPlaying = useStore((s) => s.setPlaying);
+  const stepRef = useRef(0);
+  const currentStep = useStore((s) => s.currentStep);
+
+  useEffect(() => {
+    stepRef.current = currentStep;
+  }, [currentStep]);
+
+  useEffect(() => {
+    if (!playing) return;
+    const interval = Math.max(40, 250 / speed);
+    const id = window.setInterval(() => {
+      const next = stepRef.current + 1;
+      if (next >= events.length) {
+        setPlaying(false);
+        return;
+      }
+      stepRef.current = next;
+      setCurrentStep(next);
+    }, interval);
+    return () => window.clearInterval(id);
+  }, [playing, speed, events.length, setCurrentStep, setPlaying]);
+}

frontend/src/lib/api.ts

@@ -0,0 +1,53 @@
+import type { Graph, LoadResponse, RunEvent } from "../types";
+
+async function jsonOrThrow<T>(res: Response): Promise<T> {
+  if (!res.ok) {
+    const t = await res.text();
+    throw new Error(t || `${res.status} ${res.statusText}`);
+  }
+  return res.json();
+}
+
+export async function loadModel(model_id: string): Promise<LoadResponse> {
+  const res = await fetch("/api/load", {
+    method: "POST",
+    headers: { "content-type": "application/json" },
+    body: JSON.stringify({ model_id }),
+  });
+  return jsonOrThrow<LoadResponse>(res);
+}
+
+export async function fetchGraph(): Promise<Graph> {
+  const res = await fetch("/api/graph");
+  return jsonOrThrow<Graph>(res);
+}
+
+export async function runModel(prompt: string | null): Promise<{ run_id: string; steps: number; elapsed_s: number }> {
+  const res = await fetch("/api/run", {
+    method: "POST",
+    headers: { "content-type": "application/json" },
+    body: JSON.stringify({ prompt }),
+  });
+  return jsonOrThrow(res);
+}
+
+export function streamRun(run_id: string, onEvent: (e: RunEvent) => void, onEnd: () => void, onError: (e: Event) => void): WebSocket {
+  const proto = location.protocol === "https:" ? "wss:" : "ws:";
+  const ws = new WebSocket(`${proto}//${location.host}/ws/run/${run_id}`);
+  ws.onmessage = (m) => {
+    try {
+      const data = JSON.parse(m.data);
+      if (data.kind === "stream_end") {
+        onEnd();
+        ws.close();
+        return;
+      }
+      onEvent(data as RunEvent);
+    } catch {
+      // ignore malformed
+    }
+  };
+  ws.onerror = onError;
+  ws.onclose = onEnd;
+  return ws;
+}

frontend/src/lib/colors.ts

@@ -0,0 +1,38 @@
+export const KIND_COLORS: Record<string, { bg: string; border: string; text: string; accent: string }> = {
+  Embedding:  { bg: "#1e293b", border: "#38bdf8", text: "#e0f2fe", accent: "#38bdf8" },
+  Attention:  { bg: "#1e1b4b", border: "#a78bfa", text: "#ede9fe", accent: "#a78bfa" },
+  MLP:        { bg: "#172554", border: "#60a5fa", text: "#dbeafe", accent: "#60a5fa" },
+  Norm:       { bg: "#1f2937", border: "#9ca3af", text: "#f3f4f6", accent: "#9ca3af" },
+  Linear:     { bg: "#0f172a", border: "#22d3ee", text: "#cffafe", accent: "#22d3ee" },
+  Conv:       { bg: "#0c4a6e", border: "#fbbf24", text: "#fef3c7", accent: "#fbbf24" },
+  Activation: { bg: "#3f3f46", border: "#facc15", text: "#fef9c3", accent: "#facc15" },
+  Dropout:    { bg: "#27272a", border: "#a1a1aa", text: "#fafafa", accent: "#a1a1aa" },
+  Pooler:     { bg: "#1e3a8a", border: "#34d399", text: "#d1fae5", accent: "#34d399" },
+  Block:      { bg: "#3b0764", border: "#f472b6", text: "#fce7f3", accent: "#f472b6" },
+  Head:       { bg: "#7c2d12", border: "#fb923c", text: "#ffedd5", accent: "#fb923c" },
+  Other:      { bg: "#262626", border: "#737373", text: "#e5e5e5", accent: "#737373" },
+};
+
+export function colorFor(kind: string) {
+  return KIND_COLORS[kind] ?? KIND_COLORS.Other;
+}
+
+export function viridis(t: number): string {
+  // Approximate viridis colormap. t in [0,1].
+  const stops: [number, number, number][] = [
+    [68, 1, 84],
+    [59, 82, 139],
+    [33, 145, 140],
+    [94, 201, 98],
+    [253, 231, 37],
+  ];
+  const x = Math.max(0, Math.min(1, t)) * (stops.length - 1);
+  const i = Math.floor(x);
+  const f = x - i;
+  const a = stops[i];
+  const b = stops[Math.min(i + 1, stops.length - 1)];
+  const r = Math.round(a[0] + (b[0] - a[0]) * f);
+  const g = Math.round(a[1] + (b[1] - a[1]) * f);
+  const bl = Math.round(a[2] + (b[2] - a[2]) * f);
+  return `rgb(${r},${g},${bl})`;
+}

frontend/src/lib/layout.ts

@@ -0,0 +1,64 @@
+import dagre from "@dagrejs/dagre";
+import type { Edge, Node } from "reactflow";
+import type { Graph } from "../types";
+
+const NODE_W = 220;
+const NODE_H = 64;
+
+export type LayoutOptions = {
+  granularity: number; // 0 = roots only, 1 = blocks, 2 = leaves
+};
+
+const DEPTH_FOR_GRAN: Record<number, number> = { 0: 1, 1: 3, 2: 99 };
+
+export function layoutGraph(graph: Graph, opts: LayoutOptions): { nodes: Node[]; edges: Edge[] } {
+  const maxDepth = DEPTH_FOR_GRAN[opts.granularity] ?? 3;
+  const visible = new Set<string>();
+  for (const n of graph.nodes) {
+    if (n.id === "<root>") continue;
+    if (n.depth <= maxDepth) visible.add(n.id);
+  }
+  // If no leaves at this depth, surface their nearest visible ancestors instead.
+  if (visible.size === 0) {
+    for (const n of graph.nodes) if (n.id !== "<root>") visible.add(n.id);
+  }
+
+  const idToNode = new Map(graph.nodes.map((n) => [n.id, n]));
+  const g = new dagre.graphlib.Graph();
+  g.setDefaultEdgeLabel(() => ({}));
+  g.setGraph({ rankdir: "TB", nodesep: 30, ranksep: 50 });
+  visible.forEach((id) => g.setNode(id, { width: NODE_W, height: NODE_H }));
+
+  // Build parent→child edges only between visible nodes.
+  const edges: Edge[] = [];
+  for (const e of graph.edges) {
+    if (e.from === "<root>") continue;
+    if (visible.has(e.from) && visible.has(e.to)) {
+      g.setEdge(e.from, e.to);
+      edges.push({
+        id: `${e.from}->${e.to}`,
+        source: e.from,
+        target: e.to,
+        type: "smoothstep",
+        animated: false,
+        style: { stroke: "#6b7280", strokeWidth: 1.5 },
+      });
+    }
+  }
+
+  dagre.layout(g);
+
+  const nodes: Node[] = [];
+  visible.forEach((id) => {
+    const layout = g.node(id);
+    const data = idToNode.get(id);
+    if (!data || !layout) return;
+    nodes.push({
+      id,
+      type: "module",
+      position: { x: layout.x - NODE_W / 2, y: layout.y - NODE_H / 2 },
+      data: { ...data },
+    });
+  });
+  return { nodes, edges };
+}

frontend/src/main.tsx

@@ -0,0 +1,11 @@
+import React from "react";
+import ReactDOM from "react-dom/client";
+import App from "./App";
+import "reactflow/dist/style.css";
+import "./styles.css";
+
+ReactDOM.createRoot(document.getElementById("root")!).render(
+  <React.StrictMode>
+    <App />
+  </React.StrictMode>,
+);

frontend/src/store.ts

@@ -0,0 +1,93 @@
+import { create } from "zustand";
+import type { Graph, LoadResponse, RunEvent, Snapshot } from "./types";
+
+type SnapshotMap = Record<string, Snapshot>;
+
+type State = {
+  modelId: string;
+  loadInfo: LoadResponse | null;
+  graph: Graph | null;
+  events: RunEvent[];
+  snapshotsByPath: SnapshotMap;
+  attentionByLayer: Record<number, any>;
+  inputUI: any | null;
+  topK: any[] | null;
+  currentStep: number;
+  playing: boolean;
+  speed: number;
+  selectedPath: string | null;
+  loading: boolean;
+  running: boolean;
+  error: string | null;
+  granularity: number; // 0 top-level, 1 blocks, 2 leaves
+
+  setModelId: (s: string) => void;
+  setLoading: (b: boolean) => void;
+  setRunning: (b: boolean) => void;
+  setError: (s: string | null) => void;
+  setLoadInfo: (l: LoadResponse | null) => void;
+  setGraph: (g: Graph | null) => void;
+  resetRun: () => void;
+  appendEvent: (ev: RunEvent) => void;
+  setCurrentStep: (n: number) => void;
+  setPlaying: (b: boolean) => void;
+  setSpeed: (n: number) => void;
+  setSelectedPath: (p: string | null) => void;
+  setGranularity: (n: number) => void;
+};
+
+export const useStore = create<State>((set) => ({
+  modelId: "prajjwal1/bert-tiny",
+  loadInfo: null,
+  graph: null,
+  events: [],
+  snapshotsByPath: {},
+  attentionByLayer: {},
+  inputUI: null,
+  topK: null,
+  currentStep: 0,
+  playing: false,
+  speed: 1,
+  selectedPath: null,
+  loading: false,
+  running: false,
+  error: null,
+  granularity: 1,
+
+  setModelId: (s) => set({ modelId: s }),
+  setLoading: (b) => set({ loading: b }),
+  setRunning: (b) => set({ running: b }),
+  setError: (s) => set({ error: s }),
+  setLoadInfo: (l) => set({ loadInfo: l }),
+  setGraph: (g) => set({ graph: g }),
+  resetRun: () =>
+    set({
+      events: [],
+      snapshotsByPath: {},
+      attentionByLayer: {},
+      inputUI: null,
+      topK: null,
+      currentStep: 0,
+      playing: false,
+    }),
+  appendEvent: (ev) =>
+    set((s) => {
+      const next: Partial<State> = { events: [...s.events, ev] };
+      if (ev.kind === "module" && ev.payload?.snapshot) {
+        next.snapshotsByPath = { ...s.snapshotsByPath, [ev.path]: ev.payload.snapshot };
+      } else if (ev.kind === "attention") {
+        next.attentionByLayer = { ...s.attentionByLayer, [ev.payload.layer]: ev.payload.snapshot };
+      } else if (ev.kind === "input") {
+        next.inputUI = ev.payload;
+      } else if (ev.kind === "logits") {
+        next.topK = ev.payload.top_k;
+        next.snapshotsByPath = { ...s.snapshotsByPath, [ev.path]: ev.payload.snapshot };
+      }
+      return next as State;
+    }),
+  setCurrentStep: (n) => set({ currentStep: n }),
+  setPlaying: (b) => set({ playing: b }),
+  setSpeed: (n) => set({ speed: n }),
+  setSelectedPath: (p) => set({ selectedPath: p }),
+  setGranularity: (n) => set({ granularity: n }),
+}));

frontend/src/styles.css

@@ -0,0 +1,68 @@
+* { box-sizing: border-box; }
+html, body, #root { height: 100%; margin: 0; }
+body {
+  font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, sans-serif;
+  background: #0b1220;
+  color: #e2e8f0;
+}
+
+button { font-family: inherit; }
+
+.btn {
+  background: #1e293b;
+  color: #e2e8f0;
+  border: 1px solid #334155;
+  border-radius: 6px;
+  padding: 5px 12px;
+  font-size: 12px;
+  cursor: pointer;
+  transition: background 120ms, border-color 120ms;
+}
+.btn:hover:not(:disabled) { background: #334155; border-color: #475569; }
+.btn:disabled { opacity: 0.5; cursor: not-allowed; }
+
+.btn-primary {
+  background: #38bdf8;
+  color: #0b1220;
+  border-color: #0ea5e9;
+  font-weight: 600;
+}
+.btn-primary:hover:not(:disabled) { background: #0ea5e9; }
+
+.tab {
+  background: transparent;
+  color: #94a3b8;
+  border: 1px solid #334155;
+  border-radius: 4px;
+  padding: 3px 8px;
+  font-size: 11px;
+  cursor: pointer;
+  transition: background 120ms, color 120ms;
+}
+.tab:hover { color: #e2e8f0; }
+.tab-active { background: #38bdf8; color: #0b1220; border-color: #0ea5e9; font-weight: 600; }
+
+.pill {
+  background: #1e293b;
+  color: #cbd5e1;
+  border: 1px solid #334155;
+  border-radius: 999px;
+  padding: 2px 8px;
+  font-size: 10px;
+  font-family: ui-monospace, monospace;
+}
+
+input[type="number"], input[type="text"] {
+  background: #0b1220;
+  color: #e2e8f0;
+  border: 1px solid #334155;
+  border-radius: 4px;
+  padding: 3px 6px;
+  font-family: ui-monospace, monospace;
+  font-size: 11px;
+}
+
+::-webkit-scrollbar { width: 8px; height: 8px; }
+::-webkit-scrollbar-track { background: #0b1220; }
+::-webkit-scrollbar-thumb { background: #334155; border-radius: 4px; }
+::-webkit-scrollbar-thumb:hover { background: #475569; }

frontend/src/types.ts

@@ -0,0 +1,57 @@
+export type Heatmap = {
+  shape: [number, number];
+  values: number[][];
+  vmin: number;
+  vmax: number;
+};
+
+export type Snapshot = {
+  shape: number[];
+  dtype: string;
+  stats: { mean: number; std: number; min: number; max: number; abs_max: number };
+  sample: number[];
+  heatmap: Heatmap | null;
+};
+
+export type GraphNode = {
+  id: string;
+  path: string;
+  module_class: string;
+  kind: string;
+  params: number;
+  params_own: number;
+  config: Record<string, unknown>;
+  children: string[];
+  depth: number;
+  is_leaf: boolean;
+};
+
+export type GraphEdge = { from: string; to: string; kind: string };
+
+export type Graph = { root: string; nodes: GraphNode[]; edges: GraphEdge[] };
+
+export type RunEventKind = "input" | "module" | "attention" | "logits" | "done" | "stream_end";
+
+export type RunEvent = {
+  step: number;
+  path: string;
+  module_class: string;
+  kind: RunEventKind;
+  payload: any;
+};
+
+export type LoadResponse = {
+  model_id: string;
+  modality: "text" | "vision" | "audio";
+  head_kind: string;
+  param_count: number;
+  config: Record<string, any>;
+  tokenizer: null | {
+    name_or_path: string;
+    vocab_size: number;
+    model_max_length: number;
+    special_tokens: string[];
+    is_fast: boolean;
+  };
+  processor: string | null;
+};

frontend/src/vite-env.d.ts

@@ -0,0 +1,2 @@
+/// <reference types="vite/client" />
+declare module "gif.js";

frontend/tsconfig.json

@@ -0,0 +1,20 @@
+{
+  "compilerOptions": {
+    "target": "ES2020",
+    "useDefineForClassFields": true,
+    "lib": ["ES2020", "DOM", "DOM.Iterable"],
+    "module": "ESNext",
+    "skipLibCheck": true,
+    "moduleResolution": "bundler",
+    "allowImportingTsExtensions": true,
+    "resolveJsonModule": true,
+    "isolatedModules": true,
+    "noEmit": true,
+    "jsx": "react-jsx",
+    "strict": true,
+    "noUnusedLocals": false,
+    "noUnusedParameters": false,
+    "noFallthroughCasesInSwitch": true
+  },
+  "include": ["src"]
+}

frontend/vite.config.ts

@@ -0,0 +1,17 @@
+import { defineConfig } from "vite";
+import react from "@vitejs/plugin-react";
+
+export default defineConfig({
+  plugins: [react()],
+  server: {
+    port: 5173,
+    proxy: {
+      "/api": "http://localhost:7860",
+      "/ws":  { target: "ws://localhost:7860", ws: true },
+    },
+  },
+  build: {
+    outDir: "dist",
+    sourcemap: false,
+  },
+});

requirements.txt

@@ -0,0 +1,13 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.4.1
+transformers==4.46.3
+tokenizers>=0.20
+safetensors>=0.4
+huggingface_hub>=0.26
+fastapi==0.115.5
+uvicorn[standard]==0.32.1
+pydantic>=2.8
+numpy>=1.26
+pillow>=10.4
+sentencepiece>=0.2
+protobuf>=4.25