Update README.md

README.md

@@ -41,8 +41,8 @@ This model is still in its infancy and it's meant to be constantly updated and t
   - #### Stable:
     - None
   - #### Beta:
-    - [vae.9710000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.9710000.pt): Trained from scratch for 9.71M steps with **dim: 32**,**vq_codebook_dim: 8192** and **vq_codebook_size: 8192**.
-    - [maskgit.4345000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/maskgit.4345000.pt): Maskgit trained from the beta VAE for 4.34M steps.
+    - [vae.12145000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.12145000.pt): Trained from scratch for 12.14M steps with **dim: 32**,**vq_codebook_dim: 8192** and **vq_codebook_size: 8192**.
+    - [maskgit.5125000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/maskgit.5125000.pt): Maskgit trained from the beta VAE for 5.12M steps.
 
   Note: Checkpoints under the Beta section are updated daily or at least 3-4 times a week. While the beta checkpoints can be used as they are, only the latest version is kept on the repo and the older checkpoints are removed when a new one
   is uploaded to keep the repo clean.
@@ -57,22 +57,24 @@ The model was trained on the following dataset:
 - **Hardware:** 1 x Nvidia RTX 3050 GPU
 - **Hours Trained:** NaN.
 - **Gradient Accumulations**: 10
-- **Batch:** 1
+- **Batch Size:** 1
 - **Learning Rate:** 1e-5
-- **Learning Rate Scheduler:** `cosine_with_restarts`
+- **Learning Rate Scheduler:** `constant_with_warmup`
 - **Scheduler Power:** 1.0
 - **Optimizer:** Adam
 - **Warmup Steps:** 10,000
-- **Number of Cycles:** 1
+- **Number of Cycles:** 200
 - **Resolution/Image Size**: First trained at a resolution of 64x64, then increased to 256x256 and then to 512x512. Check the notes down below for more details on this.
 - **Dimension:** 32
 - **vq_codebook_dim:** 8192
 - **vq_codebook_size:** 8192
+- **num_tokens:** 8192
+- **seq_len:** 1024
 - **heads:** 8
 - **depth:** 4
 - **Random Crop:** True
-- **Total MaskGit Training Steps:** 4,345,000
-- **Total VAE Training Steps:** 9,710,000
+- **Total MaskGit Training Steps:** 5,125,000
+- **Total VAE Training Steps:** 12,145,000
 
   Note: On Muse we can change the image_size or resolution at any time without having to train the model from scratch again, this allows us to first train the model at low resolution using the same `dim` and `vq_codebook_size` to train faster and then we can increase the `image_size` and use a higher resolution once the model has trained enough.