update readme

README.md

@@ -1,32 +1,31 @@
-# Neural Cellular Automata (Based on https://distill.pub/2020/growing-ca/) implemented in Jax (Flax)
-
-![Gecko gif](https://raw.githubusercontent.com/shyamsn97/jax-nca/main/images/gecko.gif?token=GHSAT0AAAAAABTB4G7FLAJSLDHSIOQONS3IYTB5ZEA)
-
+---
+tags:
+- image-generation
 ---
 
+# Neural Cellular Automata (Based on https://distill.pub/2020/growing-ca/) implemented in Jax (Flax)
+
 
 ## Installation
-from source:
-```
+from source
+
+```bash
 git clone git@github.com:shyamsn97/jax-nca.git
 cd jax-nca
 python setup.py install
 ```
 
 from PYPI
-```
+
+```bash
 pip install jax-nca
 ```
----
 
 ## How do NCAs work?
 For more information, view the awesome article https://distill.pub/2020/growing-ca/ -- Mordvintsev, et al., "Growing Neural Cellular Automata", Distill, 2020
 
 Image below describes a single update step: https://github.com/distillpub/post--growing-ca/blob/master/public/figures/model.svg
 
-![NCA update](https://raw.githubusercontent.com/shyamsn97/jax-nca/main/images/model.svg?token=GHSAT0AAAAAABTB4G7FOWOPXEUYVLBGRNSWYTB5YUA)
-
----
 
 ## Why Jax?
 
@@ -34,7 +33,7 @@ Image below describes a single update step: https://github.com/distillpub/post--
 
 NCAs are autoregressive models like RNNs, where new states are calculated from previous ones. With jax, we can make these operations a lot more performant with `jax.lax.scan`  and `jax.jit` (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)
 
-Instead of writing the nca growth process as:
+Instead of writing the nca growth process as
 
 ```python
 def multi_step(params, nca, current_state, num_steps):
@@ -66,14 +65,13 @@ def multi_step(params, nca, current_state, num_steps):
 ```
 The actual multi_step implementation can be found here: https://github.com/shyamsn97/jax-nca/blob/main/jax_nca/nca.py#L103
 
----
 
 ## Usage
 See [notebooks/Gecko.ipynb](notebooks/Gecko.ipynb) for a full example
 
 <b> Currently there's a bug with the stochastic update, so only `cell_fire_rate = 1.0` works at the moment </b>
 
-Creating and using NCA:
+Creating and using NCA
 
 ```python
 class NCA(nn.Module):
@@ -117,7 +115,8 @@ update = nca.apply({"params":params}, nca_seed, jax.random.PRNGKey(10))
 final_state, nca_states = nca.multi_step(poarams, nca_seed, jax.random.PRNGKey(10), num_steps=32)
 ```
 
-To train the NCA:
+To train the NCA
+
 ```python
 from jax_nca.dataset import ImageDataset
 from jax_nca.trainer import EmojiTrainer