diff --git a/README.md b/README.md
index 242afb67..803fcd20 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # DNA Diffusion
 
-<img src='https://github.com/pinellolab/DNA-Diffusion/raw/1413ec9c974d54a1e2885a250540dc89955d10c0/docs/images/diff_first.gif'> </img>
+<img src='https://raw.githubusercontent.com/pinellolab/DNA-Diffusion/main/docs/images/diff_first_lossy.gif?inline=true'> </img>
 
 <p align="center">
     <em>Generative modeling of regulatory DNA sequences with diffusion probabilistic models.</em>
diff --git a/docs/images/diff_first_lossy.gif b/docs/images/diff_first_lossy.gif
new file mode 100644
index 00000000..75cee53f
Binary files /dev/null and b/docs/images/diff_first_lossy.gif differ
diff --git a/docs/reference/dnadiffusion.md b/docs/reference/dnadiffusion.md
index 62a3a955..ac125b9a 100644
--- a/docs/reference/dnadiffusion.md
+++ b/docs/reference/dnadiffusion.md
@@ -1,6 +1,47 @@
 # Reference
 
-::: dnadiffusion
+## dnadiffusion
+
+::: dnadiffusion.trainer
+    options:
+      show_root_heading: true
+      show_source: false
+
+::: dnadiffusion.configs
+    options:
+      show_root_heading: true
+      show_source: false
+
+::: dnadiffusion.sample
+    options:
+      show_root_heading: true
+      show_source: false
+
+## data
+
+::: dnadiffusion.data.dataloader
+    options:
+      show_root_heading: true
+      show_source: false
+
+## models
+
+::: dnadiffusion.models.modules
+    options:
+      show_root_heading: true
+      show_source: false
+
+::: dnadiffusion.models.training_modules
+    options:
+      show_root_heading: true
+      show_source: false
+    
+::: dnadiffusion.models.networks
+    options:
+      show_root_heading: true
+      show_source: false
+
+::: dnadiffusion.models.unet
     options:
-      show_root_heading: false
+      show_root_heading: true
       show_source: false