Update README.md

README.md

@@ -49,6 +49,7 @@ model-index:
 
 [![arXiv](https://img.shields.io/badge/PDF-arXiv-blue)](https://www.arxiv.org/abs/2503.03965)
 [![Code](https://img.shields.io/badge/Code-GitHub-blue)](https://github.com/facebookresearch/all-atom-diffusion-transformer/)
+[![Weights](https://img.shields.io/badge/Weights-HuggingFace-blue)](https://huggingface.co/chaitjo/all-atom-diffusion-transformer)
 [![X](https://img.shields.io/badge/X_thread-@chaitjo-blue)](https://x.com/chaitjo/status/1899114667219304525)
 [![Slides](https://img.shields.io/badge/Slides-chaitjo.com-blue)](https://www.chaitjo.com/publication/joshi-2025-allatom/All_Atom_Diffusion_Transformers_Slides.pdf)
 
@@ -58,13 +59,12 @@ Independent reproduction of the paper [*"All-atom Diffusion Transformers: Unifie
 
 Note that these checkpoints are the result of an independent reproduction of this research by Chaitanya K. Joshi, and may not correspond to the exact models/performance metrics reported in the final manuscript.
 
-These checkpoints can be used to run inference as described in the [README on GitHub](https://github.com/facebookresearch/all-atom-diffusion-transformer/). Alternatively, here is a minimal script for loading an ADiT checkpoint:
-```python
-from src.models.ldm_module import LatentDiffusionLitModule
-ckpt_path = "path-to-ckpt-file"
-model = LatentDiffusionLitModule.load_from_checkpoint(ckpt_path)
-model.eval()  # for sampling and inference
-```
+These checkpoints can be used to run inference as described in the [README on GitHub](https://github.com/facebookresearch/all-atom-diffusion-transformer/). 
+
+Here is a minimal notebook for loading an ADiT checkpoint and sampling some crystals or molecules:
+<a target="_blank" href="https://colab.research.google.com/drive/1wHXsP0SHZ-Lx6Brgg-osuvTFrWw3M7oW?usp=sharing">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
 
 Examples of 10,000 sampled crystals and molecules are also available:
 - [Crystals as CIF files](https://huggingface.co/chaitjo/all-atom-diffusion-transformer/resolve/main/ADiT_crystals_mp20.zip)