Add model card

README.md

@@ -0,0 +1,58 @@
+---
+license: mit
+tags:
+  - diffusion
+  - ddpm
+  - sparse-autoencoder
+  - jet-physics
+  - quark-gluon
+---
+
+# JetDDPM-SAE — Checkpoints
+
+Two stage model on quark/gluon calorimeter jet images (3-channel, 125×125):
+1. DDPM UNet trained to generate realistic jet images.
+2. Sparse Autoencoder on frozen UNet `ups[0]` activations for interpretability.
+
+## Checkpoints
+
+| File | Description |
+|------|-------------|
+| `diff_checkpoint.pth` | Trained diffusion UNet (raw `state_dict`) |
+| `sae_upblock_checkpoint.pth` | SAE trained on ups[0] activations at t=100 |
+
+## Loading diffusion model
+
+```python
+import torch
+from src.diffusion import Unet
+from src.noise_scheduler import LinNoiseScheduler
+
+model = Unet()
+model.load_state_dict(torch.load("diff_checkpoint.pth", map_location="cpu"))
+model.eval()
+
+scheduler = LinNoiseScheduler(num_of_timesteps=1000)
+# generate: start from noise and call scheduler.sample_prev_timestep()
+```
+
+## Loading SAE
+
+```python
+import torch
+from src.sae import SparseAutoencoder
+
+ckpt  = torch.load("sae_upblock_checkpoint.pth", map_location="cpu")
+model = SparseAutoencoder(input_dim=ckpt["input_dim"], hidden_dim=1024)
+model.load_state_dict(ckpt["model"])
+model.eval()
+
+# pass global-avg-pooled ups[0] activations → sparse feature vector z
+recon, z = model(activation_vector)
+```
+
+## Training details
+
+- Diffusion: MSE noise prediction, Adam lr=1e-4, 50 epochs, batch 32
+- SAE hook: `ups[0]` (first decoder block, 31×31, 256ch), fixed t=100
+- SAE: MSE + L1(z)·1e-3, Adam lr=1e-3, 50 epochs on 30k activations